Merge pull request #424 from Xilinx/kosh.pow.to.mul.canonicalization.for.qdq

Enable Pow to Mul canonicalization for quantized exponents

Author: kosh-rai

src/Dialect/ONNX/ONNXOps/OpHelper.cpp

@@ -596,8 +596,8 @@ bool isDenseONNXConstant(Value result) {
 template <typename RESULT_TYPE>
 RESULT_TYPE getScalarValue(ElementsAttr denseAttr, Type type) {
   Type elementaryType = getElementTypeOrSelf(type);
-  if (elementaryType.isInteger(16) || elementaryType.isInteger(32) ||
-      elementaryType.isInteger(64)) {
+  if (elementaryType.isInteger(8) || elementaryType.isInteger(16) ||
+      elementaryType.isInteger(32) || elementaryType.isInteger(64)) {
     auto valueIt = denseAttr.getValues<IntegerAttr>().begin();
     return static_cast<RESULT_TYPE>(mlir::cast<IntegerAttr>(*valueIt).getInt());
   } else if (mlir::isa<FloatType>(elementaryType)) {
@@ -794,6 +794,38 @@ IgnoreDiagnostic::~IgnoreDiagnostic() {
 
 bool hasIntegerPowerExponent(ONNXPowOp *op, int64_t &exponentValue) {
   Value exponent = op->getY();
+  // In case of QDQ quantized models: If exponent is from a DequantizeLinear op,
+  // we want to check the dequantized value of the exponent
+  if (auto dequantizeOp = mlir::dyn_cast_or_null<ONNXDequantizeLinearOp>(
+          exponent.getDefiningOp())) {
+    ElementsAttr xAttr = getElementAttributeFromONNXValue(dequantizeOp.getX());
+    ElementsAttr scaleAttr =
+        getElementAttributeFromONNXValue(dequantizeOp.getXScale());
+    ElementsAttr zeroPointAttr =
+        getElementAttributeFromONNXValue(dequantizeOp.getXZeroPoint());
+
+    if (!(isScalarConstantTensor(dequantizeOp.getXScale()) &&
+            isScalarConstantTensor(dequantizeOp.getXZeroPoint())))
+      return false;
+
+    auto x = getScalarValue<double>(xAttr, xAttr.getElementType());
+    auto scale = getScalarValue<double>(scaleAttr, scaleAttr.getElementType());
+    auto zeroPoint =
+        getScalarValue<double>(zeroPointAttr, zeroPointAttr.getElementType());
+
+    // Calculate dequantized value for exponent (This is an approximation and
+    // isn't expected to match the actual calculation done by the
+    // DequantizeLinear op. However, it should be good enough for checking that
+    // the exponent is an integer)
+    double dequantizedExponent = (x - zeroPoint) * scale;
+
+    if (dequantizedExponent == ceil(dequantizedExponent)) {
+      exponentValue = static_cast<int64_t>(dequantizedExponent);
+      return true;
+    }
+    return false;
+  }
+
   ElementsAttr elementAttr = getElementAttributeFromONNXValue(exponent);
   if (!elementAttr)
     return false;

test/mlir/onnx/onnx_canonicalization.mlir

@@ -1387,6 +1387,30 @@ func.func @expand_pow_into_constant(%arg0: tensor<3x4x5xf32>) -> tensor<3x4x5xf3
 // CHECK:           onnx.Return [[VAR_0_]] : tensor<3x4x5xf32>
 // CHECK:         }
 }
+// -----
+
+func.func @test_pow_into_mul_with_qdq(%arg0: tensor<1x3x80x80x2xi8>) -> tensor<1x3x80x80x2xi8> {
+    %0 = onnx.Constant dense<2.500000e-01> : tensor<f32>
+    %1 = onnx.Constant dense<3.125000e-02> : tensor<f32>
+    %2 = onnx.Constant dense<64> : tensor<i8>
+    %3 = onnx.Constant dense<0> : tensor<i8>
+    %6 = "onnx.DequantizeLinear"(%arg0, %0, %3) {axis = 1 : si64, block_size = 0 : si64} : (tensor<1x3x80x80x2xi8>, tensor<f32>, tensor<i8>) -> tensor<1x3x80x80x2xf32>
+    %7 = "onnx.DequantizeLinear"(%2, %1, %3) {axis = 1 : si64, block_size = 0 : si64} : (tensor<i8>, tensor<f32>, tensor<i8>) -> tensor<f32>
+    %8 = "onnx.Pow"(%6, %7) : (tensor<1x3x80x80x2xf32>, tensor<f32>) -> tensor<1x3x80x80x2xf32>
+    %9 = "onnx.QuantizeLinear"(%8, %1, %3) {axis = 1 : si64, block_size = 0 : si64, output_dtype = 0 : si64, saturate = 1 : si64} : (tensor<1x3x80x80x2xf32>, tensor<f32>, tensor<i8>) -> tensor<1x3x80x80x2xi8>
+    return %9 : tensor<1x3x80x80x2xi8>
+
+// CHECK-LABEL: func.func @test_pow_into_mul_with_qdq
+// CHECK-SAME: ([[PARAM_0_:%.+]]: tensor<1x3x80x80x2xi8>) -> tensor<1x3x80x80x2xi8> {
+// CHECK: [[VAR_0_:%.+]] = onnx.Constant dense<2.500000e-01> : tensor<f32>
+// CHECK: [[VAR_1_:%.+]] = onnx.Constant dense<3.125000e-02> : tensor<f32>
+// CHECK: [[VAR_2_:%.+]] = onnx.Constant dense<0> : tensor<i8>
+// CHECK: [[VAR_3_:%.+]] = "onnx.DequantizeLinear"([[PARAM_0_]], [[VAR_0_]], [[VAR_2_]]) {axis = 1 : si64, block_size = 0 : si64} : (tensor<1x3x80x80x2xi8>, tensor<f32>, tensor<i8>) -> tensor<1x3x80x80x2xf32>
+// CHECK: [[VAR_4_:%.+]] = "onnx.Mul"([[VAR_3_]], [[VAR_3_]]) : (tensor<1x3x80x80x2xf32>, tensor<1x3x80x80x2xf32>) -> tensor<1x3x80x80x2xf32>
+// CHECK: [[VAR_5_:%.+]] = "onnx.QuantizeLinear"([[VAR_4_]], [[VAR_1_]], [[VAR_2_]]) {axis = 1 : si64, block_size = 0 : si64, output_dtype = 0 : si64, saturate = 1 : si64} : (tensor<1x3x80x80x2xf32>, tensor<f32>, tensor<i8>) -> tensor<1x3x80x80x2xi8>
+// CHECK: return [[VAR_5_]] : tensor<1x3x80x80x2xi8>
+// CHECK: }
+}
 
 // -----