Extend Quant dialect with Quantile Quantization type (#53)

* Expanding Quant dialect with Quantile Quantized type

* Adding quantile mlir tests

* Adding check on quantiles array size and updated mlir tests

Author: sartil

mlir/include/mlir/Dialect/Quant/IR/QuantDialectBytecode.td

@@ -81,6 +81,35 @@ def UniformQuantizedPerAxisType: DialectType<(type
   }];
 }
 
+def QuantileQuantizedType: DialectType<(type
+  VarInt:$flags,
+  Type:$storageType,
+  Type:$expressedType,
+  Array<DoubleAPFloatList>:$quantiles,
+  DoubleAPFloat:$scale,
+  SignedVarInt:$zeroPoint,
+  SignedVarInt:$storageTypeMin,
+  SignedVarInt:$storageTypeMax
+)>;
+
+def QuantileQuantizedPerAxisType: DialectType<(type
+  VarInt:$flags,
+  Type:$storageType,
+  Type:$expressedType,
+  VarInt:$quantizedDimension,
+  SignedVarInt:$storageTypeMin,
+  SignedVarInt:$storageTypeMax,
+  Array<DoubleAPFloatList>:$quantiles,
+  Array<DoubleAPFloatList>:$scales,
+  Array<SignedVarIntList>:$zeroPoints
+)> {
+  // Note: builder order differs from bytecode.
+  let cBuilder = [{
+      get<$_resultType>(context, flags, storageType, expressedType, quantiles, scales,
+        zeroPoints, quantizedDimension, storageTypeMin, storageTypeMax)
+  }];
+}
+
 /// This enum contains marker codes used to indicate which attribute is
 /// currently being decoded, and how it should be decoded. The order of these
 /// codes should generally be unchanged, as any changes will inevitably break
@@ -93,7 +122,9 @@ def QuantDialectTypes : DialectTypes<"Quant"> {
     AnyQuantizedTypeWithExpressedType,
     CalibratedQuantizedType,
     UniformQuantizedType,
-    UniformQuantizedPerAxisType
+    UniformQuantizedPerAxisType,
+    QuantileQuantizedType,
+    QuantileQuantizedPerAxisType
   ];
 }
 

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

@@ -25,6 +25,8 @@ struct QuantizedTypeStorage;
 struct AnyQuantizedTypeStorage;
 struct UniformQuantizedTypeStorage;
 struct UniformQuantizedPerAxisTypeStorage;
+struct QuantileQuantizedTypeStorage;
+struct QuantileQuantizedPerAxisTypeStorage;
 struct CalibratedQuantizedTypeStorage;
 
 } // namespace detail
@@ -394,6 +396,128 @@ class UniformQuantizedPerAxisType
   }
 };
 
+/// QuantileQuantizedType derives from UniformQuantizedType and adds to it a
+/// look up table array of quantile values.
+///
+/// Syntax synopsis:
+///   Per-layer, all parameters expressed:
+///     !quant<quantile[StorageType:ExpressedType]{Quantiles}:{Scale:ZeroPoint}>
+///   Per-layer, optional parameters omitted:
+///     !quant<quantile[StorageType]{Quantiles}:{Scale}>
+///
+///   StorageType: 'i'|'u' NumBits
+///   ExpressedType: 'f16', 'f32', 'bf16', 'f64'
+///   Quantiles: Quantile+
+///   Quantile: A legal double value
+///   Scale: A legal double value
+///   ZeroPoint: An integer value
+class QuantileQuantizedType
+    : public Type::TypeBase<QuantileQuantizedType, UniformQuantizedType,
+                            detail::QuantileQuantizedTypeStorage> {
+public:
+  using Base::Base;
+  using Base::getChecked;
+
+  static constexpr StringLiteral name = "quant.quantile";
+
+  /// Gets an instance of the type with all parameters specified but not
+  /// checked.
+  static QuantileQuantizedType get(unsigned flags, Type storageType,
+                                   Type expressedType,
+                                   ArrayRef<double> quantiles, double scale,
+                                   int64_t zeroPoint, int64_t storageTypeMin,
+                                   int64_t storageTypeMax);
+
+  static QuantileQuantizedType
+  getChecked(function_ref<InFlightDiagnostic()> emitError, unsigned flags,
+             Type storageType, Type expressedType, ArrayRef<double> quantiles,
+             double scale, int64_t zeroPoint, int64_t storageTypeMin,
+             int64_t storageTypeMax);
+
+  /// Verifies construction invariants and issues errors/warnings.
+  static LogicalResult verifyInvariants(function_ref<InFlightDiagnostic()> emitError,
+                              unsigned flags, Type storageType,
+                              Type expressedType, ArrayRef<double> quantiles,
+                              double scale, int64_t zeroPoint,
+                              int64_t storageTypeMin, int64_t storageTypeMax);
+
+  /// Gets the quantile values
+  ArrayRef<double> getQuantiles() 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; }
+};
+
+/// Represents per-axis QuantileQuantizedType (also known as per-channel
+/// quantization).
+///
+/// Syntax synopsis:
+///   Per-axis, all parameters expressed:
+///     !quant<quantile[StorageType:ExpressedType:QuantizedDim]{Quantiles}:{QuantParams}>
+///   Per-axis, optional parameters omitted:
+///     !quant<quantile[StorageType]{Quantiles}:{Scale}>
+///
+///   StorageType: 'i'|'u' NumBits
+///   ExpressedType: 'f16', 'f32', 'bf16', 'f64'
+///   QuantizedDim: An integer value
+///   Quantiles: Quantile+
+///   Quantile: A legal double value
+///   QuantParams: (Scale ':' ZeroPoint)+
+///   Scale: A legal double value
+///   ZeroPoint: An integer value
+class QuantileQuantizedPerAxisType
+    : public Type::TypeBase<QuantileQuantizedPerAxisType,
+                            UniformQuantizedPerAxisType,
+                            detail::QuantileQuantizedPerAxisTypeStorage> {
+public:
+  using Base::Base;
+  using Base::getChecked;
+
+  static constexpr StringLiteral name = "quant.quantile_per_axis";
+
+  /// Gets an instance of the type with all parameters specified but not
+  /// checked.
+  static QuantileQuantizedPerAxisType
+  get(unsigned flags, Type storageType, Type expressedType,
+      ArrayRef<double> quantiles, ArrayRef<double> scales,
+      ArrayRef<int64_t> zeroPoints, int32_t quantizedDimension,
+      int64_t storageTypeMin, int64_t storageTypeMax);
+
+  /// Gets an instance of the type with all specified parameters checked.
+  /// Returns a nullptr convertible type on failure.
+  static QuantileQuantizedPerAxisType
+  getChecked(function_ref<InFlightDiagnostic()> emitError, unsigned flags,
+             Type storageType, Type expressedType, ArrayRef<double> quantiles,
+             ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
+             int32_t quantizedDimension, int64_t storageTypeMin,
+             int64_t storageTypeMax);
+
+  /// Verifies construction invariants and issues errors/warnings.
+  static LogicalResult verifyInvariants(function_ref<InFlightDiagnostic()> emitError,
+                              unsigned flags, Type storageType,
+                              Type expressedType, ArrayRef<double> quantiles,
+                              ArrayRef<double> scales,
+                              ArrayRef<int64_t> zeroPoints,
+                              int32_t quantizedDimension,
+                              int64_t storageTypeMin, int64_t storageTypeMax);
+
+  /// Gets the quantile values
+  ArrayRef<double> getQuantiles() 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() && !llvm::is_contained(getZeroPoints(), 0);
+  }
+};
+
 /// A quantized type that infers its range from given min/max values.
 ///
 /// Typical syntax:

mlir/include/mlir/Dialect/Quant/QuantOpsBase.td

@@ -0,0 +1,84 @@
+//===- QuantOpsBase.td - Quantization dialect base ---------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Predicates for types in the Quantization dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef DIALECT_QUANT_QUANT_OPS_BASE_
+#define DIALECT_QUANT_QUANT_OPS_BASE_
+
+include "mlir/IR/OpBase.td"
+
+def Quantization_Dialect : Dialect {
+  let name = "quant";
+  let cppNamespace = "::mlir::quant";
+
+  let useDefaultTypePrinterParser = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Quantization type definitions
+//===----------------------------------------------------------------------===//
+
+class quant_TypedPrimitiveOrContainer<Type etype> :
+    Type<Or<[etype.predicate,
+                TensorOf<[etype]>.predicate,
+                VectorOf<[etype]>.predicate]>,
+         "primitive/tensor/vector of " # etype.summary>;
+
+// An implementation of QuantizedType.
+def quant_QuantizedType :
+    Type<CPred<"::llvm::isa<mlir::quant::QuantizedType>($_self)">, "QuantizedType">;
+
+// A primitive type that can represent a real value. This is either a
+// floating point value or a quantized type.
+def quant_RealPrimitiveType :
+    Type<Or<[AnyFloat.predicate, quant_QuantizedType.predicate]>,
+    "real valued primitive (float or quantized type)">;
+
+// A primitive type that can represent a storage value. This is either an
+// integer or quantized type.
+def quant_StoragePrimitiveType :
+    Type<Or<[AnySignlessInteger.predicate, quant_QuantizedType.predicate]>,
+    "quantized storage primitive (integer or quantized type)">;
+
+// A primitive or container of RealPrimitiveType.
+def quant_RealValueType :
+    quant_TypedPrimitiveOrContainer<quant_RealPrimitiveType>;
+
+// A primitive or container of StoragePrimitiveType.
+def quant_StorageValueType :
+    quant_TypedPrimitiveOrContainer<quant_StoragePrimitiveType>;
+
+// Either a real valued or storage primitive or container type.
+def quant_RealOrStorageValueType :
+    Type<Or<[quant_RealValueType.predicate, quant_StorageValueType.predicate]>,
+    "real valued or storage primitive or container type">;
+
+// An implementation of UniformQuantizedType.
+def quant_UniformQuantizedType :
+    DialectType<Quantization_Dialect,
+                CPred<"::llvm::isa<UniformQuantizedType>($_self)">,
+                "UniformQuantizedType">;
+
+// An implementation of QuantileQuantizedType.
+def quant_QuantileQuantizedType :
+    DialectType<Quantization_Dialect,
+                CPred<"::llvm::isa<QuantileQuantizedType>($_self)">,
+                "QuantileQuantizedType">;
+
+// Predicate for detecting a container or primitive of UniformQuantizedType.
+def quant_UniformQuantizedValueType :
+    quant_TypedPrimitiveOrContainer<quant_UniformQuantizedType>;
+
+// Predicate for detecting a container or primitive of QuantileQuantizedType.
+def quant_QuantileQuantizedValueType :
+    quant_TypedPrimitiveOrContainer<quant_QuantileQuantizedType>;
+
+#endif // DIALECT_QUANT_QUANT_OPS_BASE_

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

@@ -93,7 +93,8 @@ LogicalResult verifyQuantizationOp(Operation *op, QuantizedType quantizedType,
 
 void QuantDialect::initialize() {
   addTypes<AnyQuantizedType, CalibratedQuantizedType, UniformQuantizedType,
-           UniformQuantizedPerAxisType>();
+           UniformQuantizedPerAxisType, QuantileQuantizedType,
+           QuantileQuantizedPerAxisType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Quant/IR/QuantOps.cpp.inc"