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_convert_fake_quant_op.cc b/tao_compiler/mlir/disc/transforms/disc_convert_fake_quant_op.cc
index 301221329b4..e240c656783 100644
--- a/tao_compiler/mlir/disc/transforms/disc_convert_fake_quant_op.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_convert_fake_quant_op.cc
@@ -143,11 +143,122 @@ struct QuantizedConvLikeOpConverter
   }
 };
 
+Value buildQuantizedDotGeneralOp(mhlo::DotOp dotOp, PatternRewriter& rewriter,
+                                 Location loc, Type outTy, Value input,
+                                 Value weight,
+                                 mhlo_disc::FakeQuantOp inputFakeQuantOp,
+                                 mhlo_disc::FakeQuantOp weightFakeQuantOp,
+                                 mhlo_disc::FakeQuantOp resultFakeQuantOp) {
+  SmallVector<int64_t> lhsContractingDims{1};
+  SmallVector<int64_t> rhsContractingDims{0};
+  auto dotDimensionAttr = mhlo::DotDimensionNumbersAttr::get(
+      rewriter.getContext(), {}, {}, lhsContractingDims, rhsContractingDims);
+  return rewriter.create<mhlo_disc::QuantizedDotGeneralOp>(
+      loc, outTy, input, weight, inputFakeQuantOp.scale(),
+      inputFakeQuantOp.zero_point(), weightFakeQuantOp.scale(),
+      weightFakeQuantOp.zero_point(), resultFakeQuantOp.scale(),
+      resultFakeQuantOp.zero_point(), dotDimensionAttr,
+      resultFakeQuantOp.use_symmetric(), weightFakeQuantOp.axis(),
+      resultFakeQuantOp.use_dynamic());
+}
+
+Value buildQuantizedDotGeneralOp(mhlo::DotGeneralOp dotOp,
+                                 PatternRewriter& rewriter, Location loc,
+                                 Type outTy, Value input, Value weight,
+                                 mhlo_disc::FakeQuantOp inputFakeQuantOp,
+                                 mhlo_disc::FakeQuantOp weightFakeQuantOp,
+                                 mhlo_disc::FakeQuantOp resultFakeQuantOp) {
+  return rewriter.create<mhlo_disc::QuantizedDotGeneralOp>(
+      loc, outTy, input, weight, inputFakeQuantOp.scale(),
+      inputFakeQuantOp.zero_point(), weightFakeQuantOp.scale(),
+      weightFakeQuantOp.zero_point(), resultFakeQuantOp.scale(),
+      resultFakeQuantOp.zero_point(), dotOp.dot_dimension_numbers(),
+      resultFakeQuantOp.use_symmetric(), weightFakeQuantOp.axis(),
+      resultFakeQuantOp.use_dynamic());
+}
+
+template <typename OpTy>
+struct QuantizedDotLikeOpConverter
+    : public OpRewritePattern<mhlo_disc::FakeQuantOp> {
+  using OpRewritePattern<mhlo_disc::FakeQuantOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo_disc::FakeQuantOp op,
+                                PatternRewriter& rewriter) const override {
+    auto dotOp = op.input().template getDefiningOp<OpTy>();
+    if (!dotOp) return failure();
+
+    if (isa<mhlo::DotOp>(dotOp.getOperation())) {
+      auto inputTy =
+          dotOp.lhs().getType().template dyn_cast<RankedTensorType>();
+      auto weightTy =
+          dotOp.rhs().getType().template dyn_cast<RankedTensorType>();
+      if (!inputTy || !weightTy || inputTy.getRank() != weightTy.getRank() ||
+          inputTy.getRank() != 2) {
+        return rewriter.notifyMatchFailure(
+            op, "not supported quantized gemv a.t.m.");
+      }
+    }
+
+    auto inputFakeQuantOp =
+        dotOp.lhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
+    auto weightFakeQuantOp =
+        dotOp.rhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
+    if (!inputFakeQuantOp || !weightFakeQuantOp) return failure();
+    if (inputFakeQuantOp.use_signed() != weightFakeQuantOp.use_signed() ||
+        inputFakeQuantOp.num_bits() != weightFakeQuantOp.num_bits()) {
+      return failure();
+    }
+    if (inputFakeQuantOp.use_dynamic() != op.use_dynamic() ||
+        weightFakeQuantOp.use_dynamic() != op.use_dynamic())
+      return failure();
+
+    Location loc = op.getLoc();
+    auto buildQuantizedTensorType = [&](mhlo_disc::FakeQuantOp fakeQuantOp) {
+      Type quantizedElemType =
+          fakeQuantOp.use_signed()
+              ? IntegerType::get(rewriter.getContext(), fakeQuantOp.num_bits())
+              : IntegerType::get(
+                    rewriter.getContext(), fakeQuantOp.num_bits(),
+                    mlir::IntegerType::SignednessSemantics::Unsigned);
+      auto inputTy = fakeQuantOp.input().getType().cast<RankedTensorType>();
+      return RankedTensorType::get(inputTy.getShape(), quantizedElemType,
+                                   inputTy.getEncoding());
+    };
+
+    // Builds a real quantized op by extracting metadata info from the
+    // fake_quant_op
+    auto buildQuantizedOpFromFakeQuantOp =
+        [&](mhlo_disc::FakeQuantOp fakeQuantOp) {
+          auto outTy = buildQuantizedTensorType(fakeQuantOp);
+          Value quantizedInput = rewriter.create<mhlo_disc::QuantizeOp>(
+              loc, outTy, fakeQuantOp.input(), fakeQuantOp.scale(),
+              fakeQuantOp.zero_point(), fakeQuantOp.use_symmetric(),
+              fakeQuantOp.axis(), fakeQuantOp.quant_min(),
+              fakeQuantOp.quant_max(), fakeQuantOp.use_dynamic());
+          return quantizedInput;
+        };
+
+    Value quantizedInput = buildQuantizedOpFromFakeQuantOp(inputFakeQuantOp);
+    Value quantizedWeight = buildQuantizedOpFromFakeQuantOp(weightFakeQuantOp);
+    Value quantizedDot = buildQuantizedDotGeneralOp(
+        dotOp, rewriter, loc, buildQuantizedTensorType(op), quantizedInput,
+        quantizedWeight, inputFakeQuantOp, weightFakeQuantOp, op);
+    Value dequantizedOut = rewriter.create<mhlo_disc::DequantizeOp>(
+        loc, op.getType(), quantizedDot, op.scale(), op.zero_point(),
+        op.use_symmetric(), op.axis(), op.use_dynamic());
+
+    rewriter.replaceOp(op, dequantizedOut);
+    return success();
+  }
+};
+
 void populateQuantizedPatterns(RewritePatternSet& patterns) {
   // clang-format off
   patterns.insert<
     QuantizedConvLikeOpConverter<mhlo::ConvolutionOp>,
-    QuantizedConvLikeOpConverter<mhlo::DynamicConvOp>
+    QuantizedConvLikeOpConverter<mhlo::DynamicConvOp>,
+    QuantizedDotLikeOpConverter<mhlo::DotOp>,
+    QuantizedDotLikeOpConverter<mhlo::DotGeneralOp>
   >(patterns.getContext());
   // clang-format on
 }


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-convert-fake-quant-op.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-convert-fake-quant-op.mlir
index 18a3fee60f0..7356d2b2035 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-convert-fake-quant-op.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-convert-fake-quant-op.mlir
@@ -138,4 +138,129 @@ func.func @fake_quant_conv(
 
   // CHECK: return %[[FAKE_QUANT_RESULT]] : tensor<?x8x6x16xf32>
   return %fake_quant_result : tensor<?x8x6x16xf32>
-}
\ No newline at end of file
+}
+
+// -----
+
+// CHECK-LABEL: @fake_quant_dot
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x?xf32>, %[[WEIGHT:.*]]: tensor<?x?xf32>
+func.func @fake_quant_dot(
+  %input: tensor<?x?xf32>, %weight: tensor<?x?xf32>,
+  %input_scale : tensor<f32>, %input_zero_point : tensor<i32>,
+  %weight_scale : tensor<?xf32>, %weight_zero_point : tensor<?xi32>,
+  %result_scale : tensor<f32>, %result_zero_point : tensor<i32>
+) -> tensor<?x?xf32> {
+  // CHECK: %[[FAKE_QUANT_INPUT:.*]] = "mhlo_disc.quantize"
+  // CHECK-SAME: %[[INPUT]]
+  %fake_quant_input = "mhlo_disc.fake_quant"(%input, %input_scale, %input_zero_point) {
+      use_signed = true,
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      num_bits = 8,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?xf32>
+
+  // CHECK: %[[FAKE_QUANT_WEIGHT:.*]] = "mhlo_disc.quantize"
+  // CHECK-SAME: %[[WEIGHT]]
+  %fake_quant_weight = "mhlo_disc.fake_quant"(%weight, %weight_scale, %weight_zero_point) {
+      use_signed = true,
+      use_symmetric = true,
+      axis = dense<[1]> : tensor<1xi64>,
+      num_bits = 8,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xf32>
+
+  // CHECK: %[[RESULT:.*]] = "mhlo_disc.quantized_dot_general"
+  // CHECK-SAME: %[[FAKE_QUANT_INPUT]]
+  // CHECK-SAME: %[[FAKE_QUANT_WEIGHT]]
+  // CHECK-SAME: lhs_contracting_dimensions = [1]
+  // CHECK-SAME: rhs_contracting_dimensions = [0]
+  %result = "mhlo.dot"(%fake_quant_input, %fake_quant_weight)
+   : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+
+  // CHECK: %[[FAKE_QUANT_RESULT:.*]] = "mhlo_disc.dequantize"
+  // CHECK-SAME: %[[RESULT]]
+  %fake_quant_result = "mhlo_disc.fake_quant"(%result, %result_scale, %result_zero_point) {
+      use_signed = true,
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      num_bits = 8,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?xf32>
+
+  // CHECK: return %[[FAKE_QUANT_RESULT]]
+  return %fake_quant_result : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fake_quant_dot_general
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x?x?x?xf32>, %[[WEIGHT:.*]]: tensor<?x?x?x?xf32>
+func.func @fake_quant_dot_general(
+  %input: tensor<?x?x?x?xf32>, %weight: tensor<?x?x?x?xf32>,
+  %input_scale : tensor<f32>, %input_zero_point : tensor<i32>,
+  %weight_scale : tensor<?xf32>, %weight_zero_point : tensor<?xi32>,
+  %result_scale : tensor<f32>, %result_zero_point : tensor<i32>
+) -> tensor<?x?x?x?xf32> {
+  // CHECK: %[[FAKE_QUANT_INPUT:.*]] = "mhlo_disc.quantize"
+  // CHECK-SAME: %[[INPUT]]
+  %fake_quant_input = "mhlo_disc.fake_quant"(%input, %input_scale, %input_zero_point) {
+      use_signed = true,
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      num_bits = 8,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x?x?x?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?x?x?xf32>
+
+  // CHECK: %[[FAKE_QUANT_WEIGHT:.*]] = "mhlo_disc.quantize"
+  // CHECK-SAME: %[[WEIGHT]]
+  %fake_quant_weight = "mhlo_disc.fake_quant"(%weight, %weight_scale, %weight_zero_point) {
+      use_signed = true,
+      use_symmetric = true,
+      axis = dense<[3]> : tensor<1xi64>,
+      num_bits = 8,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x?x?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?x?x?xf32>
+
+  // CHECK: %[[RESULT:.*]] = "mhlo_disc.quantized_dot_general"
+  // CHECK-SAME: %[[FAKE_QUANT_INPUT]]
+  // CHECK-SAME: %[[FAKE_QUANT_WEIGHT]]
+  // CHECK-SAME: lhs_batching_dimensions = [0, 1]
+  // CHECK-SAME: rhs_batching_dimensions = [0, 1]
+  // CHECK-SAME: lhs_contracting_dimensions = [3]
+  // CHECK-SAME: rhs_contracting_dimensions = [3]
+  %result = "mhlo.dot_general"(%fake_quant_input, %fake_quant_weight) {
+    dot_dimension_numbers = #mhlo.dot<
+      lhs_batching_dimensions = [0, 1],
+      rhs_batching_dimensions = [0, 1],
+      lhs_contracting_dimensions = [3],
+      rhs_contracting_dimensions = [3]
+    >}
+   : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
+
+  // CHECK: %[[FAKE_QUANT_RESULT:.*]] = "mhlo_disc.dequantize"
+  // CHECK-SAME: %[[RESULT]]
+  %fake_quant_result = "mhlo_disc.fake_quant"(%result, %result_scale, %result_zero_point) {
+      use_signed = true,
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      num_bits = 8,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x?x?x?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?x?x?xf32>
+
+  // CHECK: return %[[FAKE_QUANT_RESULT]]
+  return %fake_quant_result : tensor<?x?x?x?xf32>
+}
+