[ONNX] Add per-channel quantization support for DequantizeLinear op (#4090)

Fixes nod-ai/SHARK-ModelDev#914.

---------

Signed-off-by: Vivek Khandelwal <[email protected]>

Author: vivekkhandelwal1

lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp

@@ -2250,66 +2250,91 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
         Value zeropoint = operands[2];
 
         auto operandTy = cast<Torch::ValueTensorType>(operand.getType());
-
-        auto operandETy = operandTy.getDtype();
         auto scaleTy = dyn_cast<Torch::ValueTensorType>(scale.getType());
         if (!scaleTy || !scaleTy.hasSizes())
           return rewriter.notifyMatchFailure(binder.op, "requires known rank");
         if (!resultType.hasDtype())
           return rewriter.notifyMatchFailure(binder.op,
                                              "requires known result dtype");
 
-        bool rank0 = scaleTy.getSizes().size() == 0;
-        bool length1 =
-            scaleTy.getSizes().size() == 1 && scaleTy.getSizes()[0] == 1;
-
-        if (!rank0 && !length1)
-          return rewriter.notifyMatchFailure(binder.op,
-                                             "unimplemented: non-scalar scale");
+        int64_t scaleRank = scaleTy.getSizes().size();
+        if (scaleRank > 1)
+          return rewriter.notifyMatchFailure(
+              binder.op, "unimplemented: only per-tensor or per-axis "
+                         "quantization supported");
         auto qTensorTy = getQTorchTypeFromTorchIntType(operandTy);
         if (!qTensorTy) {
           return rewriter.notifyMatchFailure(binder.op,
                                              "unsupported result dtype");
         }
 
-        scale = rewriter.create<Torch::AtenItemOp>(
-            loc, rewriter.getType<Torch::FloatType>(), scale);
-
+        auto operandETy = operandTy.getDtype();
         bool fpOperand = isa<mlir::FloatType>(operandETy);
-        Type zeropointTy = rewriter.getType<Torch::IntType>();
-        if (fpOperand)
-          zeropointTy = rewriter.getType<Torch::FloatType>();
-
-        zeropoint =
-            rewriter.create<Torch::AtenItemOp>(loc, zeropointTy, zeropoint);
-
-        if (fpOperand) {
-          Value none = rewriter.create<Torch::ConstantNoneOp>(loc);
-          Value cstFalse = rewriter.create<Torch::ConstantBoolOp>(loc, false);
-          auto tyVal = Torch::getScalarTypeForType(resultType.getDtype());
-          Value tyConst = rewriter.create<Torch::ConstantIntOp>(
-              loc, rewriter.getType<Torch::IntType>(),
-              rewriter.getIntegerAttr(rewriter.getIntegerType(64),
-                                      static_cast<int64_t>(tyVal)));
-          Value toDtype = rewriter.create<Torch::AtenToDtypeOp>(
-              loc, resultType, operand, tyConst,
-              /*non_blocking=*/cstFalse, /*copy=*/cstFalse,
-              /*memory_format=*/none);
-
-          Value one = rewriter.create<Torch::ConstantFloatOp>(
-              loc, rewriter.getF64FloatAttr(1.0));
-          Value sub = rewriter.create<Torch::AtenSubScalarOp>(
-              loc, resultType, toDtype, zeropoint, one);
-          rewriter.replaceOpWithNewOp<Torch::AtenMulScalarOp>(
-              binder.op, resultType, sub, scale);
+        bool isPerTensorQuantization = false;
+        if (scaleRank == 0 ||
+            llvm::all_of(scaleTy.getSizes(), [](int64_t s) { return s == 1; }))
+          isPerTensorQuantization = true;
+
+        // (TODO) Case: Per-Channel Quantization for floating point input.
+        if (scaleRank == 1 && fpOperand)
+          return rewriter.notifyMatchFailure(
+              binder.op, "unimplemented: support for per-Channel Quantization "
+                         "for floating point input not present");
+
+        if (isPerTensorQuantization) {
+          scale = rewriter.create<Torch::AtenItemOp>(
+              loc, rewriter.getType<Torch::FloatType>(), scale);
+
+          Type zeropointTy = rewriter.getType<Torch::IntType>();
+          if (fpOperand)
+            zeropointTy = rewriter.getType<Torch::FloatType>();
+          zeropoint =
+              rewriter.create<Torch::AtenItemOp>(loc, zeropointTy, zeropoint);
+        }
+
+        if (!fpOperand) {
+          Value quantize;
+          // Case 1: Per-Tensor Quantization for non-floating point input.
+          if (isPerTensorQuantization) {
+            quantize =
+                rewriter.create<Torch::Aten_MakePerTensorQuantizedTensorOp>(
+                    loc, qTensorTy, operand, scale, zeropoint);
+          } else {
+            // Case 2: Per-Channel Quantization for non-floating point input.
+            int64_t axis;
+            if (binder.s64IntegerAttr(axis, "axis", 1))
+              return failure();
+
+            Value cstAxis = rewriter.create<Torch::ConstantIntOp>(
+                loc, rewriter.getI64IntegerAttr(axis));
+            quantize =
+                rewriter.create<Torch::Aten_MakePerChannelQuantizedTensorOp>(
+                    loc, qTensorTy, operand, scale, zeropoint, cstAxis);
+          }
+          rewriter.replaceOpWithNewOp<Torch::AtenDequantizeSelfOp>(
+              binder.op, resultType, quantize);
           return success();
         }
 
-        auto quantize =
-            rewriter.create<Torch::Aten_MakePerTensorQuantizedTensorOp>(
-                loc, qTensorTy, operand, scale, zeropoint);
-        rewriter.replaceOpWithNewOp<Torch::AtenDequantizeSelfOp>(
-            binder.op, resultType, quantize);
+        // Case 3: Per-Tensor Quantization for floating point input.
+        Value none = rewriter.create<Torch::ConstantNoneOp>(loc);
+        Value cstFalse = rewriter.create<Torch::ConstantBoolOp>(loc, false);
+        auto tyVal = Torch::getScalarTypeForType(resultType.getDtype());
+        Value tyConst = rewriter.create<Torch::ConstantIntOp>(
+            loc, rewriter.getType<Torch::IntType>(),
+            rewriter.getIntegerAttr(rewriter.getIntegerType(64),
+                                    static_cast<int64_t>(tyVal)));
+        Value toDtype = rewriter.create<Torch::AtenToDtypeOp>(
+            loc, resultType, operand, tyConst,
+            /*non_blocking=*/cstFalse, /*copy=*/cstFalse,
+            /*memory_format=*/none);
+
+        Value one = rewriter.create<Torch::ConstantFloatOp>(
+            loc, rewriter.getF64FloatAttr(1.0));
+        Value sub = rewriter.create<Torch::AtenSubScalarOp>(
+            loc, resultType, toDtype, zeropoint, one);
+        rewriter.replaceOpWithNewOp<Torch::AtenMulScalarOp>(
+            binder.op, resultType, sub, scale);
         return success();
       });
   patterns.onOp("Div", 7,

test/Conversion/TorchOnnxToTorch/simple_ops_a_to_f.mlir

@@ -817,6 +817,18 @@ func.func @test_dequantizelinear_fp8(%arg0: !torch.vtensor<[6],f8E4M3FN>, %arg1:
 
 // -----
 
+// CHECK-LABEL: @test_dequantizelinear_per_channel_si8
+func.func @test_dequantizelinear_per_channel_si8(%arg0: !torch.vtensor<[64,3,3,3],si8>, %arg1: !torch.vtensor<[64],f32>, %arg2: !torch.vtensor<[64],si8>) -> !torch.vtensor<[64,3,3,3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 21 : si64} {
+  %0 = torch.operator "onnx.DequantizeLinear"(%arg0, %arg1, %arg2) {torch.onnx.axis = 0 : si64} : (!torch.vtensor<[64,3,3,3],si8>, !torch.vtensor<[64],f32>, !torch.vtensor<[64],si8>) -> !torch.vtensor<[64,3,3,3],f32>
+  // CHECK: %[[AXIS:.+]] = torch.constant.int 0
+  // CHECK: %[[MAKE:.+]] = torch.aten._make_per_channel_quantized_tensor %arg0, %arg1, %arg2,
+  // CHECK: %[[DEQ:.+]] = torch.aten.dequantize.self %[[MAKE]]
+  // CHECK: return %[[DEQ]]
+  return %0 : !torch.vtensor<[64,3,3,3],f32>
+}
+
+// -----
+
 // CHECK-LABEL: @test_div_bcast
 func.func @test_div_bcast(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[5],f32>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 14 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
   // CHECK: torch.aten.div.Tensor %arg0, %arg1 : !torch.vtensor<[3,4,5],f32>, !torch.vtensor<[5],f32> -> !torch.vtensor<[3,4,5],f32>