Add scale and zero_point to xten_nn.quantize/dequantize

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

@@ -113,59 +113,77 @@ 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) }];
 
   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) }];
+  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

@@ -90,37 +90,49 @@ 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());
+    const auto inputType =
+        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())));
+    const llvm::APFloat scale = quantizeOp.getScale();
+    const llvm::APFloat inverseScale =
+        llvm::APFloat::getOne(scale.getSemantics()) / scale;
 
-    // Create a constant that represents the (1 / 2 ^ shift)
-    RankedTensorType constType =
-        createSplatType(inputType.getRank(), rewriter.getF32Type());
+    const RankedTensorType constType =
+        createSplatType(inputType.getRank(), inputType.getElementType());
     auto constOp = rewriter.create<tosa::ConstOp>(
         quantizeOp->getLoc(), constType,
-        DenseFPElementsAttr::get(constType, {scale}));
+        DenseFPElementsAttr::get(constType, {inverseScale}));
 
-    // 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 +168,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 +185,36 @@ 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());
+
+    const llvm::APFloat scale = dequantizeOp.getScale();
 
     // 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();
   }

lib/Dialect/XTenNN/IR/XTenNNOps.cpp

@@ -35,6 +35,7 @@
 #include "xten/Dialect/XTenNN/IR/XTenNNBase.h"
 #include "xten/Dialect/XTenNN/Interfaces/EnclaveOpInterfaces.h"
 
+#include <cfenv>
 #include <cstdint>
 
 using namespace mlir;
@@ -444,25 +445,141 @@ LogicalResult SubgraphOp::inferReturnTypeComponents(
 // XTenNNDialect
 //===----------------------------------------------------------------------===//
 
+namespace {
+std::optional<int32_t> getShiftValue(float constValue) {
+  const float log2Value = std::log2f(constValue);
+
+  // The log2 of the value must not have fractions.
+  if (std::roundf(log2Value) != log2Value)
+    return {};
+
+  return static_cast<int32_t>(log2Value);
+}
+
+float getScaleFromShift(int32_t shift) {
+  std::feclearexcept(FE_ALL_EXCEPT);
+  errno = 0;
+  const float scale = exp2f(shift);
+  assert(!std::fetestexcept(FE_OVERFLOW));
+  assert(errno == 0);
+  return scale;
+}
+} // namespace
+
+void amd::xten_nn::QuantizeOp::build(mlir::OpBuilder &odsBuilder,
+                                     mlir::OperationState &odsState,
+                                     mlir::Type output, mlir::Value input,
+                                     mlir::FloatAttr scale,
+                                     mlir::IntegerAttr zeroPoint) {
+  const bool zeroPointIsZero = zeroPoint.getValue().isZero();
+  assert(scale.getType().isF32());
+  const auto shiftValue = getShiftValue(scale.getValue().convertToFloat());
+  if (zeroPointIsZero && shiftValue) {
+    return build(odsBuilder, odsState, output, input,
+                 odsBuilder.getSI32IntegerAttr(*shiftValue), scale, zeroPoint);
+  }
+  return build(odsBuilder, odsState, output, input, nullptr, scale, zeroPoint);
+}
+
+void amd::xten_nn::QuantizeOp::build(mlir::OpBuilder &odsBuilder,
+                                     mlir::OperationState &odsState,
+                                     mlir::Type output, mlir::Value input,
+                                     int32_t shift) {
+  const auto outputElemType = cast<TensorType>(output).getElementType();
+  return build(odsBuilder, odsState, output, input,
+               odsBuilder.getSI32IntegerAttr(shift),
+               odsBuilder.getF32FloatAttr(getScaleFromShift(shift)),
+               odsBuilder.getIntegerAttr(outputElemType, 0));
+}
+
+LogicalResult amd::xten_nn::QuantizeOp::verify() {
+  if (getResult().getType().getElementType() != getZeroPointAttr().getType()) {
+    return emitOpError("Result elem type needs to match match zero point type");
+  }
+  // if shift is set, zero point needs to be zero and scale needs to match
+  if (getShift()) {
+    const auto computedShift = getShiftValue(getScale().convertToFloat());
+    if (!computedShift || computedShift != *getShift()) {
+      return emitOpError(
+          "Shift set, but does not match shift calculated from scale");
+    }
+    if (!getZeroPoint().isZero()) {
+      return emitOpError("Shift set, but zero_point not zero");
+    }
+  }
+
+  return success();
+}
+
 OpFoldResult amd::xten_nn::QuantizeOp::fold(FoldAdaptor adaptor) {
-  // Fold away cases where a xten_nn.quantize is preceeded by xten_nn.dequantize
-  // that uses the same shift factor and has same types.
+  // Fold away cases where a xten_nn.quantize is preceeded by
+  // xten_nn.dequantize that uses the same scale factor, zeroPoint and has same
+  // types.
 
   auto dequantizeOp =
       dyn_cast_or_null<amd::xten_nn::DequantizeOp>(getInput().getDefiningOp());
   if (!dequantizeOp)
     return {};
 
-  if (!dequantizeOp->hasOneUse() || dequantizeOp.getShift() != getShift())
-    return {};
-
   auto dequantizeInput = dequantizeOp.getInput();
   if (dequantizeInput.getType() != getType())
     return {};
 
+  if (!dequantizeOp->hasOneUse() || dequantizeOp.getScale() != getScale() ||
+      dequantizeOp.getZeroPoint() != getZeroPoint())
+    return {};
+
   return dequantizeInput;
 }
 
+void amd::xten_nn::DequantizeOp::build(mlir::OpBuilder &odsBuilder,
+                                       mlir::OperationState &odsState,
+                                       mlir::Type output, mlir::Value input,
+                                       mlir::FloatAttr scale,
+                                       mlir::IntegerAttr zeroPoint) {
+  const bool zeroPointIsZero = zeroPoint.getValue().isZero();
+  assert(scale.getType().isF32());
+  const auto shiftValue = getShiftValue(scale.getValue().convertToFloat());
+  if (zeroPointIsZero && shiftValue) {
+    return build(odsBuilder, odsState, output, input,
+                 odsBuilder.getSI32IntegerAttr(*shiftValue), scale, zeroPoint);
+  }
+  return build(odsBuilder, odsState, output, input, nullptr, scale, zeroPoint);
+}
+
+void amd::xten_nn::DequantizeOp::build(mlir::OpBuilder &odsBuilder,
+                                       mlir::OperationState &odsState,
+                                       mlir::Type output, mlir::Value input,
+                                       int32_t shift) {
+  const auto inputElemType = cast<TensorType>(input.getType()).getElementType();
+  return build(odsBuilder, odsState, output, input,
+               odsBuilder.getSI32IntegerAttr(shift),
+               odsBuilder.getF32FloatAttr(getScaleFromShift(shift)),
+               odsBuilder.getIntegerAttr(inputElemType, 0));
+}
+
+LogicalResult amd::xten_nn::DequantizeOp::verify() {
+  // Input elem type should match zero point type
+  if (cast<TensorType>(getOperand().getType()).getElementType() !=
+      getZeroPointAttr().getType()) {
+    return emitOpError(
+        "Operand elem type needs to match match zero point type");
+  }
+  // if shift is set, zero point needs to be zero and scale needs to match
+  if (getShift()) {
+    const auto computedShift = getShiftValue(getScale().convertToFloat());
+    if (!computedShift || computedShift != *getShift()) {
+      return emitOpError(
+          "Shift set, but does not match shift calculated from scale");
+    }
+    if (!getZeroPoint().isZero()) {
+      return emitOpError("Shift set, but zero_point not zero");
+    }
+  }
+
+  return success();
+}
+
 OpFoldResult amd::xten_nn::GroupQuantizeOp::fold(FoldAdaptor adaptor) {
   // Fold away cases where a xten_nn.group_quantize is preceeded by
   // xten_nn.group_dequantize that uses the same shift factor and has same