Merge pull request #148 from Xilinx/jrickert.qdq

Add scale and zero_point to xten_nn.quantize/dequantize

include/xten/Dialect/XTenNN/IR/XTenNNOps.td

@@ -113,59 +113,82 @@ def XTenNN_QuantizeOp: XTenNN_Op<"quantize", [
             Elementwise,
             Pure,
             SameOperandsAndResultShape]> {
-  let summary = "Quantizes a float32 tensor to a signless or unsigned integer tensor of given width.";
+  let summary = "Quantizes a float tensor to a signless or unsigned integer tensor of given width.";
   let description = [{
-    Quantizes a given float32 tensor into a signless or unsigned integer tensor of given width.
+    Quantizes a given float tensor into a signless or unsigned integer tensor of given width.
     Since tosa is using signless/unsigned types currently, we also consider signless integer types for
     signed ones when the type is not unsigned until tosa support signed integers.
 
     Applies the following linear quantization to the input tensor x:
-      y = round( x / 2^shift )
+      y = round((x / scale) + zero_point)
 
-    Where 2^shift is equal to the scale of the quantize operation and
-    the shift is an attribute of the operation in si32.
+    Iff log2(scale) is representable as si32 and zero_point == 0, shift is set to log2(scale).
+    In this case the quantization is equal to:
+      y = round( x / 2^shift )
 
     Round will saturate to the range of the output type and the rounding mode is set to half
     to nearest even.
   }];
 
   let arguments = (ins
-    F32Tensor:$input,
-    SI32Attr:$shift
+    XTenNN_AnyFloatTensor:$input,
+    OptionalAttr<SI32Attr>:$shift,
+    F32Attr:$scale, // Restricted to F32 for now, but may be relaxed in the future 
+    XTenNN_AnyIntegerAttr:$zero_point
   );
 
   let results = (outs XTenNN_AnySignlessOrUnsignedIntegerTensor:$output);
+  let builders = [
+    OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "::mlir::FloatAttr":$scale, "::mlir::IntegerAttr":$zeroPoint)>,
+    OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "int32_t":$shift)>
+  ];
 
-  let assemblyFormat = [{ `(`$input `:` type($input)`)` attr-dict `->` type($output) }];
 
+  //  `(`$input `:` type($input)`)` attr-dict `->` type($output);
+  // Zero point is optional if scale is set, defaults to 0
+  // If shift is set but not scale, scale is based on the shift. Setting a zero point is not allowed in this case
+  let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
+  let hasVerifier = 1;
 }
 
 def XTenNN_DequantizeOp: XTenNN_Op<"dequantize", [
             Elementwise,
             Pure,
             SameOperandsAndResultShape]> {
-  let summary = "Dequantizes a signless/unsigned integer tensor of given bitwidth to a float32 tensor.";
+  let summary = "Dequantizes a signless/unsigned integer tensor of given bitwidth to a float tensor.";
   let description = [{
-    Dequantizes a signless/unsigned integer tensor of given bitwidth to a float32 tensor.
+    Dequantizes a signless/unsigned integer tensor of given bitwidth to a float tensor.
     Since tosa is using signless/unsigned types currently, we also consider signless integer types for
     signed ones when the type is not unsigned until tosa support signed integers.
 
     Applies the following linear dequantization to the input tensor x:
-      y = x  * ( 2^shift )
+      y = (x - zero_point) * scale
 
-    Where 2^shift is equal to scale of the dequantize operation and
-    the shift is an attribute of the operation in si32.
+    Iff log2(scale) is representable as si32 and zero_point == 0, shift is set to log2(scale).
+    In this case the dequantization is equal to:
+      y = x  * ( 2^shift )
   }];
 
   let arguments = (ins
     XTenNN_AnySignlessOrUnsignedIntegerTensor:$input,
-    SI32Attr:$shift
+    OptionalAttr<SI32Attr>:$shift,
+    F32Attr:$scale, // Restricted to F32 for now, but may be relaxed in the future
+    XTenNN_AnyIntegerAttr:$zero_point
   );
 
-  let results = (outs F32Tensor:$output);
+  let results = (outs XTenNN_AnyFloatTensor:$output);
+
+  let builders = [
+    OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "::mlir::FloatAttr":$scale, "::mlir::IntegerAttr":$zeroPoint)>,
+    OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "int32_t":$shift)>
+  ];
 
-  let assemblyFormat = [{ `(`$input `:` type($input)`)` attr-dict `->` type($output) }];
+  //  `(`$input `:` type($input)`)` attr-dict `->` type($output);
+  // Zero point is optional if scale is set, defaults to 0
+  // If shift is set but not scale, scale is based on the shift. Setting a zero point is not allowed in this case
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
 }
 
 def XTenNN_GroupQuantizeOp: XTenNN_Op<"group_quantize", [

include/xten/Dialect/XTenNN/IR/XTenNNTypes.td

@@ -22,6 +22,15 @@ def XTenNN_AnySignlessOrUnsignedIntegerTensor : TensorOf<
 
 def XTenNN_AnyFloatTensor : TensorOf<[AnyFloat]>;
 
+def XTenNN_AnyIntegerAttr:
+    TypedAttrBase<
+      AnyInteger, "IntegerAttr",
+      CPred<"::llvm::isa<::mlir::IntegerAttr>($_self)">,
+      "Any integer attr"> {
+  let returnType = [{ ::llvm::APInt }];
+  let constBuilderCall = ?;
+}
+
 def XTenNN_AnyIntegerOrFloat : AnyTypeOf<[AnyInteger, AnyFloat], "Integer or Float">;
 
 #endif // XTENNN_TYPES

lib/Conversion/TosaToXTenNN.cpp

@@ -224,7 +224,11 @@ class FoldMulsToQDQOps : public OpRewritePattern<tosa::MulOp> {
     }
 
     // Sum the shifts of the quantize, dequantize and update the operations
-    llvm::APInt shiftSum(32, dequantizeOp.getShift(), true);
+    if (!dequantizeOp.getShift()) {
+      return rewriter.notifyMatchFailure(dequantizeOp.getLoc(),
+                                         "Dequantize op has no shift");
+    }
+    llvm::APInt shiftSum(32, *dequantizeOp.getShift(), true);
     bool overflow = false;
     shiftSum = shiftSum.sadd_ov(dequantizeShift, overflow);
     if (overflow) {

lib/Conversion/XTenNNToTosa.cpp

@@ -81,6 +81,22 @@ IntegerMinMax calculateMinMaxOfElementType(TensorType type) {
   return IntegerMinMax{minValue.getSExtValue(), maxValue.getSExtValue()};
 }
 
+namespace {
+APFloat convertF32AttrToFloatTy(FloatAttr attr, Type typeToConvertTo) {
+  // Convert from f32 to the float type that is actually used
+  assert(attr.getType().isF32());
+  assert(isa<FloatType>(typeToConvertTo));
+  auto floatResultType = cast<FloatType>(typeToConvertTo);
+  APFloat scale = attr.getValue();
+  bool losesInfo;
+  // Ignore inaccuracies, there is nothing we can do.
+  [[maybe_unused]] const auto conversionResult =
+      scale.convert(floatResultType.getFloatSemantics(),
+                    llvm::RoundingMode::NearestTiesToEven, &losesInfo);
+  return scale;
+}
+} // namespace
+
 class QuantizeOp : public OpRewritePattern<amd::xten_nn::QuantizeOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
@@ -90,37 +106,52 @@ class QuantizeOp : public OpRewritePattern<amd::xten_nn::QuantizeOp> {
     // The QDQ operations only work on tensors, if they are not, then the
     // verifiers should find the error. At the moment, only signless tensors
     // are supported.
-    auto outputType = cast<TensorType>(quantizeOp->getResult(0).getType());
+    const auto outputType =
+        cast<TensorType>(quantizeOp->getResult(0).getType());
     if (!outputType.getElementType().isSignlessInteger()) {
       return rewriter.notifyMatchFailure(
           quantizeOp.getLoc(),
           "only signless integer tensor types are supported.");
     }
-    auto inputType = dyn_cast<TensorType>(quantizeOp->getOperand(0).getType());
-
-    // Calculate (1 / 2 ^ shift)
-    llvm::APFloat scale(std::pow(static_cast<float>(2.0),
-                                 static_cast<float>(-quantizeOp.getShift())));
-
-    // Create a constant that represents the (1 / 2 ^ shift)
-    RankedTensorType constType =
-        createSplatType(inputType.getRank(), rewriter.getF32Type());
+    const auto inputType =
+        cast<TensorType>(quantizeOp->getOperand(0).getType());
+    const auto inputElementType = inputType.getElementType();
+
+    // Convert the scale from f32 to the float type that is actually used
+    const llvm::APFloat scale =
+        convertF32AttrToFloatTy(quantizeOp.getScaleAttr(), inputElementType);
+    const llvm::APFloat scaleReciprocal =
+        llvm::APFloat::getOne(scale.getSemantics()) / scale;
+
+    const RankedTensorType constType =
+        createSplatType(inputType.getRank(), inputElementType);
     auto constOp = rewriter.create<tosa::ConstOp>(
         quantizeOp->getLoc(), constType,
-        DenseFPElementsAttr::get(constType, {scale}));
+        DenseFPElementsAttr::get(constType, {scaleReciprocal}));
 
-    // Calculate (x / 2 ^ shift)
     auto mulOp = rewriter.create<tosa::MulOp>(
         quantizeOp.getLoc(), inputType, quantizeOp->getOperand(0),
         constOp->getResult(0), rewriter.getI8IntegerAttr(0));
 
+    const auto constAddType =
+        createSplatType(inputType.getRank(), outputType.getElementType());
+    auto constAddOp = rewriter.create<tosa::ConstOp>(
+        quantizeOp.getLoc(), constAddType,
+        DenseIntElementsAttr::get(constAddType, {quantizeOp.getZeroPoint()}));
+    auto constAddCastOp = rewriter.create<tosa::CastOp>(
+        quantizeOp.getLoc(), inputType, constAddOp.getResult());
+    auto zeroPointAdd = rewriter.create<tosa::AddOp>(
+        quantizeOp.getLoc(), inputType, mulOp.getResult(),
+        constAddCastOp.getResult());
+
     // TOSA only supports signed integers of i8, i16 or i32 here we convert our
     // si<?> to this types and add a clamp to mimic arbitrary bit width.
     TensorType newIntegerStorageType = getNewStorageType(outputType);
     // Cast from fp32 -> i<Bitwidth> where bit width is the supported storage
     // bit width. Either i8, i16 or i32
-    auto castOp = rewriter.create<tosa::CastOp>(
-        quantizeOp->getLoc(), newIntegerStorageType, mulOp->getResult(0));
+    auto castOp = rewriter.create<tosa::CastOp>(quantizeOp->getLoc(),
+                                                newIntegerStorageType,
+                                                zeroPointAdd->getResult(0));
 
     // Find the max and min of the signed integer type.
     IntegerMinMax intLimits = calculateMinMaxOfElementType(outputType);
@@ -156,7 +187,10 @@ class DequantizeOp : public OpRewritePattern<amd::xten_nn::DequantizeOp> {
     // The QDQ operations only work on tensors, if they are not, then the
     // verifiers should find the error. At the moment, only signless tensors
     // are supported.
-    auto inputType = cast<TensorType>(dequantizeOp->getOperand(0).getType());
+    const auto resultElementType =
+        dequantizeOp.getResult().getType().getElementType();
+    const auto inputType =
+        cast<TensorType>(dequantizeOp->getOperand(0).getType());
     if (!inputType.getElementType().isSignlessInteger()) {
       return rewriter.notifyMatchFailure(
           dequantizeOp.getLoc(),
@@ -170,26 +204,38 @@ class DequantizeOp : public OpRewritePattern<amd::xten_nn::DequantizeOp> {
         dequantizeOp.getLoc(), newIntegerStorageType,
         dequantizeOp->getOperand(0));
 
-    // We can then cast from i<8,16,32> -> fp32
+    // We can then cast from i<8,16,32> -> fp
     auto castOp = rewriter.create<tosa::CastOp>(
         dequantizeOp->getLoc(), dequantizeOp->getResult(0).getType(),
         unrealizedCast.getResult(0));
 
-    // Calculate the (x * 2 ^ shift) for the dequantize part
-    llvm::APFloat scale(std::pow(static_cast<float>(2.0),
-                                 static_cast<float>(dequantizeOp.getShift())));
+    // Do the zero_point sub on the float type to to avoid underflows
+    const auto constSubType =
+        createSplatType(inputType.getRank(), inputType.getElementType());
+    auto constSubOp = rewriter.create<tosa::ConstOp>(
+        dequantizeOp.getLoc(), constSubType,
+        DenseIntElementsAttr::get(constSubType, {dequantizeOp.getZeroPoint()}));
+    auto constSubCastOp = rewriter.create<tosa::CastOp>(
+        dequantizeOp.getLoc(), dequantizeOp.getResult().getType(),
+        constSubOp.getResult());
+    auto zeroPointSub = rewriter.create<tosa::SubOp>(
+        dequantizeOp.getLoc(), dequantizeOp.getResult().getType(),
+        castOp.getResult(), constSubCastOp.getResult());
+
+    // Convert the scale from f32 to the float type that is actually used
+    const llvm::APFloat scale =
+        convertF32AttrToFloatTy(dequantizeOp.getScaleAttr(), resultElementType);
 
     // Create a constant to hold the floating point scale we just calculated
-    auto constType =
-        createSplatType(inputType.getRank(), rewriter.getF32Type());
+    auto constType = createSplatType(inputType.getRank(), resultElementType);
     auto constOp = rewriter.create<tosa::ConstOp>(
         dequantizeOp->getLoc(), constType,
         DenseFPElementsAttr::get(constType, {scale}));
 
     // Replace the dequantize op with the new operations we just created.
     rewriter.replaceOpWithNewOp<tosa::MulOp>(
         dequantizeOp, dequantizeOp->getResult(0).getType(),
-        castOp->getResult(0), constOp->getResult(0),
+        zeroPointSub->getResult(0), constOp->getResult(0),
         rewriter.getI8IntegerAttr(0));
     return success();
   }