Few fixes to help ONNX import for quantized models

lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp

@@ -195,7 +195,7 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
                       binder.op, resultType, operand);
                   return success();
                 });
-  patterns.onOp("MatMul", 13,
+  patterns.onOp("MatMul", 1,
                 [](OpBinder binder, ConversionPatternRewriter &rewriter) {
                   Torch::ValueTensorType resultType;
                   Value lhs, rhs;

lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp

@@ -1426,7 +1426,7 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
                 });
 
   patterns.onOp(
-      "Transpose", 13,
+      "Transpose", 1,
       [](OpBinder binder, ConversionPatternRewriter &rewriter) {
         auto loc = binder.getLoc();
         Torch::ValueTensorType resultType;

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -152,11 +152,6 @@ LogicalResult torchScalarToTosaTensor(ConversionPatternRewriter &rewriter,
                      .value();
   } else if (auto intType = dtype.dyn_cast<mlir::IntegerType>()) {
     auto w = intType.getWidth();
-    if (w != 1 && w != 32 && w != 64)
-      return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
-        diag << "Unsupported integer type: " << intType;
-      });
-
     if (w == 1) {
       if (!isInValidRange<bool>(isFloat, doubleValue, isInt, intValue)) {
         return rewriter.notifyMatchFailure(
@@ -167,6 +162,17 @@ LogicalResult torchScalarToTosaTensor(ConversionPatternRewriter &rewriter,
                        : static_cast<bool>(intValue);
       tosaTensor =
           tosa::getConstTensor<bool>(rewriter, op, {d}, dshape).value();
+    } else if (w == 8) {
+      if (!isInValidRange<int8_t>(isFloat, doubleValue, isInt, intValue)) {
+        return rewriter.notifyMatchFailure(
+            op, "Supplied value of scalar constant exceeds limits "
+                "of destination type");
+      }
+      int8_t d = isFloat ? static_cast<int8_t>(doubleValue)
+                          : static_cast<int8_t>(intValue);
+      tosaTensor =
+          tosa::getConstTensor<int8_t>(rewriter, op, {d}, dshape).value();
+
     } else if (w == 32) {
       if (!isInValidRange<int32_t>(isFloat, doubleValue, isInt, intValue)) {
         return rewriter.notifyMatchFailure(
@@ -186,6 +192,10 @@ LogicalResult torchScalarToTosaTensor(ConversionPatternRewriter &rewriter,
       int64_t d = (isFloat ? static_cast<int64_t>(doubleValue) : intValue);
       tosaTensor =
           tosa::getConstTensor<int64_t>(rewriter, op, {d}, dshape).value();
+    } else {
+      return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
+        diag << "Unsupported integer type: " << intType;
+      });
     }
   } else {
     return rewriter.notifyMatchFailure(op, "Usupported element type");

lib/Dialect/Torch/Transforms/RecomposeComplexOps.cpp

@@ -600,6 +600,39 @@ class RecomposeRepeatInterleave : public OpRewritePattern<AtenRepeatInterleaveTe
   }
 };
 
+class RecomposeQuantizePerTensor : public OpRewritePattern<AtenIntReprOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenIntReprOp op,
+                                PatternRewriter &rewriter) const override {
+    auto quantize = op.getSelf().getDefiningOp<AtenQuantizePerTensorOp>();
+    if(!quantize) {
+      return rewriter.notifyMatchFailure(
+        op, "no quantize");
+    }
+
+    BaseTensorType type = dyn_cast<BaseTensorType>(op.getType());
+    if(!type)
+      return rewriter.notifyMatchFailure(op, "no quantize");
+
+    Location loc = op.getLoc();
+    auto div = rewriter.create<AtenDivScalarOp>(loc, quantize.getSelf().getType(), quantize.getSelf(), quantize.getScale());
+        Value none = rewriter.create<ConstantNoneOp>(loc);
+    Value cstFalse = rewriter.create<ConstantBoolOp>(loc, false);
+
+    auto cast = rewriter.create<AtenToDtypeOp>(
+        loc, op.getType(), div,
+        getDtypeIntValueForType(rewriter, loc, type.getDtype()),
+        /*nonBlocking=*/cstFalse, /*copy=*/cstFalse, /*memoryFormat=*/none);
+
+    Value one =
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(1));
+
+    rewriter.replaceOpWithNewOp<AtenAddScalarOp>(op, op.getType(), cast, quantize.getZeroPoint(), one);
+    return success();
+  }
+};
+
 } // namespace
 
 namespace {
@@ -621,6 +654,7 @@ class RecomposeComplexOpsPass
     patterns.add<RecomposeSplitTensorPrimListUnpackOp>(context);
     patterns.add<RecomposeChunkListUnpack>(context);
     patterns.add<RecomposeRepeatInterleave>(context);
+    patterns.add<RecomposeQuantizePerTensor>(context);
 
     GreedyRewriteConfig config;
     config.useTopDownTraversal = true;