Repository: alibaba/BladeDISC (language: python)

Issue:
Initial quantization support in BladeDISC
We're now planing to start support quantization in BladeDISC. The basic idea is:
- introduce a fake_quant op. The op will be used to pass quantization related info (e.g. scale/zero_point) from frontend to BladeDISC.
- The input of BladeDISC is still a FP32 graph, with some fake_quant ops to annotate which set of ops need to be quantized.
-  It's the responsibility of BladeDISC to replace the fake_quant op and the corresponding FP32 ops (e.g. Gemm/Conv) to the quantized versions. 

The task list is shown as follow:
- [x] TensorFlow/PyTorch fake_quant custom ops and their corresponding implementation in TF/PyTorch Dialect.
- [x] mhlo_disc fake_quantize/quantize/de-quantize IR definition.
- [x] mhlo_disc QCONV/QMATMUL IR defintion.
- [x] TF custom fake-quant -> disc fake quant
- [ ] Torch custom fake-quant -> disc fake quant
- [x] Fake_quantize + Conv + Fake_quantize  -> QConv pattern.
- [x] Fake_quantize + GEMM + Fake_quantize  -> QGEMM pattern.
- [x] Fake_quantize lowering.
- [x] Quantize op lowering patterns.
- [x] de-quantize op lowering patterns.
- [x] Placement rule for Qconv/Qmatmul.
- [ ] (low priority) QConv/QMatmul layout optimization.
- [x] QConv shape inference implementation.
- [x] QMatmul shape inference implementation.
- [x] QConv lowering to library call implementation.
- [x] QMatmul lowering to library call implementation.
- [x] QConv kernel and e2e UT for arm.
- [x] QMatMul kernel and e2e UT for arm.
- [x] QConv kernel and e2e UT for cuda.
- [x] QMatMul kernel and e2e UT for cuda.
- [ ] QConv kernel and e2e UT for x86.
- [ ] QMatMul kernel and e2e UT for x86.
- [ ] e2e demo model for ARM
- [ ] e2e demo model for X86
- [ ] e2e demo model for CUDA



Discussion / hints:

Fix patch:
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc b/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
index eac8015f8bb..bf9ccfc6643 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
@@ -86,9 +86,51 @@ struct QuantizeOpConverter : public OpRewritePattern<mhlo_disc::QuantizeOp> {
   }
 };
 
+struct DequantizeOpConverter
+    : public OpRewritePattern<mhlo_disc::DequantizeOp> {
+  using OpRewritePattern<mhlo_disc::DequantizeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo_disc::DequantizeOp op,
+                                PatternRewriter& rewriter) const override {
+    if (op.use_dynamic())
+      return rewriter.notifyMatchFailure(
+          op, "Not support dynamic dequantize a.t.m.");
+
+    auto outTy = op.getType().cast<RankedTensorType>();
+    auto scaleTy = op.scale().getType().cast<RankedTensorType>();
+    if (outTy.getElementType() != scaleTy.getElementType() ||
+        !outTy.getElementType().isF32())
+      return rewriter.notifyMatchFailure(
+          op, "Only support dequantize to f32 a.t.m.");
+
+    Location loc = op.getLoc();
+    auto inputTy = op.input().getType().cast<RankedTensorType>();
+    Value inputShape = rewriter.create<shape::ShapeOfOp>(loc, op.input());
+    Value bcastedScale = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
+        loc, outTy, op.scale(), inputShape, op.axis());
+    auto zeroPointTy = op.zero_point().getType().cast<RankedTensorType>();
+    auto bcastedInputOrZeroPointTy =
+        RankedTensorType::get(inputTy.getShape(), zeroPointTy.getElementType(),
+                              zeroPointTy.getEncoding());
+    Value bcastedZeroPoint = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
+        loc, bcastedInputOrZeroPointTy, op.zero_point(), inputShape, op.axis());
+    Value castedInput = rewriter.create<mhlo::ConvertOp>(
+        loc, bcastedInputOrZeroPointTy, op.input());
+
+    // output = (input - zero\_point) \times scale
+    Value t0 =
+        rewriter.create<mhlo::SubtractOp>(loc, castedInput, bcastedZeroPoint);
+    Value t1 = rewriter.create<mhlo::ConvertOp>(loc, outTy, t0);
+    Value t2 = rewriter.create<mhlo::MulOp>(loc, t1, bcastedScale);
+    rewriter.replaceOp(op, t2);
+    return success();
+  }
+};
+
 void populateQuantizeAndDequantizePatterns(RewritePatternSet& patterns) {
   // clang-format off
   patterns.insert<
+    DequantizeOpConverter,
     QuantizeOpConverter
   >(patterns.getContext());
   // clang-format on


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
index 5c8d6fe65e9..1d80c24fe50 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
@@ -193,3 +193,135 @@ func.func @quantize_unsigned_per_channel(
   } : (tensor<?x32x32x6xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x32x32x6xui8>
   return %quantized_input : tensor<?x32x32x6xui8>
 }
+
+// -----
+
+// CHECK-LABEL: @dequantize_signed_per_tensor
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xi8>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<f32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<i32>
+func.func @dequantize_signed_per_tensor(
+    %input: tensor<?x32x32x6xi8>,
+    %input_scale : tensor<f32>,
+    %input_zero_point : tensor<i32>) -> tensor<?x32x32x6xf32> {
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_SCALE]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xi8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
+  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
+  // CHECK: return %[[T2]]
+  %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xi8>, tensor<f32>, tensor<i32>) -> tensor<?x32x32x6xf32>
+  return %quantized_input : tensor<?x32x32x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dequantize_unsigned_per_tensor
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xui8>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<f32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<i32>
+func.func @dequantize_unsigned_per_tensor(
+    %input: tensor<?x32x32x6xui8>,
+    %input_scale : tensor<f32>,
+    %input_zero_point : tensor<i32>) -> tensor<?x32x32x6xf32> {
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_SCALE]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xui8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
+  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
+  // CHECK: return %[[T2]]
+  %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xui8>, tensor<f32>, tensor<i32>) -> tensor<?x32x32x6xf32>
+  return %quantized_input : tensor<?x32x32x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dequantize_signed_per_channel
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xi8>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<?xf32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<?xi32>
+func.func @dequantize_signed_per_channel(
+    %input: tensor<?x32x32x6xi8>,
+    %input_scale : tensor<?xf32>,
+    %input_zero_point : tensor<?xi32>) -> tensor<?x32x32x6xf32> {
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_SCALE]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xi8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
+  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
+  // CHECK: return %[[T2]]
+  %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[3]> : tensor<1xi64>,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xi8>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x32x32x6xf32>
+  return %quantized_input : tensor<?x32x32x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dequantize_unsigned_per_channel
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xui8>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<?xf32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<?xi32>
+func.func @dequantize_unsigned_per_channel(
+    %input: tensor<?x32x32x6xui8>,
+    %input_scale : tensor<?xf32>,
+    %input_zero_point : tensor<?xi32>) -> tensor<?x32x32x6xf32> {
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_SCALE]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xui8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
+  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
+  // CHECK: return %[[T2]]
+  %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[3]> : tensor<1xi64>,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xui8>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x32x32x6xf32>
+  return %quantized_input : tensor<?x32x32x6xf32>
+}