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 ad57f803c6d..301221329b4 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
@@ -42,19 +42,51 @@ namespace mlir {
 namespace disc_ral {
 namespace {
 
-struct QuantizedDynamicConvOpConverter
+template <typename OpTy>
+struct ConvLikeOpPadding {
+  static Value get(OpTy convOp, OpBuilder& builder) { return nullptr; }
+};
+
+template <>
+struct ConvLikeOpPadding<mhlo::DynamicConvOp> {
+  static Value get(mhlo::DynamicConvOp convOp, OpBuilder& builder) {
+    if (!convOp) return nullptr;
+    return convOp.d_padding();
+  }
+};
+
+template <>
+struct ConvLikeOpPadding<mhlo::ConvolutionOp> {
+  static Value get(mhlo::ConvolutionOp convOp, OpBuilder& builder) {
+    if (!convOp) return nullptr;
+    SmallVector<int32_t> paddingValues;
+    for (const auto& val : (*convOp.padding()).getValues<int64_t>()) {
+      paddingValues.push_back(static_cast<int32_t>(val));
+    }
+    RankedTensorType ty = RankedTensorType::get({paddingValues.size()},
+                                                builder.getIntegerType(32));
+    return builder.create<mhlo::ConstantOp>(
+        convOp.getLoc(), DenseIntElementsAttr::get(ty, paddingValues));
+  }
+};
+
+template <typename OpTy>
+struct QuantizedConvLikeOpConverter
     : public OpRewritePattern<mhlo_disc::FakeQuantOp> {
   using OpRewritePattern<mhlo_disc::FakeQuantOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(mhlo_disc::FakeQuantOp op,
                                 PatternRewriter& rewriter) const override {
-    auto convOp = op.input().getDefiningOp<mhlo::DynamicConvOp>();
+    auto convOp = op.input().template getDefiningOp<OpTy>();
     if (!convOp) return failure();
 
+    Value padding = ConvLikeOpPadding<OpTy>::get(convOp, rewriter);
+    if (!padding) return failure();
+
     auto inputFakeQuantOp =
-        convOp.lhs().getDefiningOp<mhlo_disc::FakeQuantOp>();
+        convOp.lhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
     auto weightFakeQuantOp =
-        convOp.rhs().getDefiningOp<mhlo_disc::FakeQuantOp>();
+        convOp.rhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
     if (!inputFakeQuantOp || !weightFakeQuantOp) return failure();
     if (inputFakeQuantOp.use_signed() != weightFakeQuantOp.use_signed() ||
         inputFakeQuantOp.num_bits() != weightFakeQuantOp.num_bits()) {
@@ -94,14 +126,13 @@ struct QuantizedDynamicConvOpConverter
     Value quantizedWeight = buildQuantizedOpFromFakeQuantOp(weightFakeQuantOp);
     Value quantizedConv = rewriter.create<mhlo_disc::QuantizedDynamicConvOp>(
         loc, buildQuantizedTensorType(op), quantizedInput, quantizedWeight,
-        convOp.d_padding(), inputFakeQuantOp.scale(),
-        inputFakeQuantOp.zero_point(), weightFakeQuantOp.scale(),
-        weightFakeQuantOp.zero_point(), op.scale(), op.zero_point(),
-        *convOp.window_strides(), *convOp.padding(), *convOp.lhs_dilation(),
-        *convOp.rhs_dilation(), *convOp.window_reversal(),
-        convOp.dimension_numbers(), convOp.feature_group_count(),
-        convOp.batch_group_count(), op.use_symmetric(),
-        weightFakeQuantOp.axis(), op.use_dynamic());
+        padding, inputFakeQuantOp.scale(), inputFakeQuantOp.zero_point(),
+        weightFakeQuantOp.scale(), weightFakeQuantOp.zero_point(), op.scale(),
+        op.zero_point(), *convOp.window_strides(), *convOp.padding(),
+        *convOp.lhs_dilation(), *convOp.rhs_dilation(),
+        *convOp.window_reversal(), convOp.dimension_numbers(),
+        convOp.feature_group_count(), convOp.batch_group_count(),
+        op.use_symmetric(), weightFakeQuantOp.axis(), op.use_dynamic());
 
     Value dequantizedOut = rewriter.create<mhlo_disc::DequantizeOp>(
         loc, op.getType(), quantizedConv, op.scale(), op.zero_point(),
@@ -115,7 +146,8 @@ struct QuantizedDynamicConvOpConverter
 void populateQuantizedPatterns(RewritePatternSet& patterns) {
   // clang-format off
   patterns.insert<
-    QuantizedDynamicConvOpConverter
+    QuantizedConvLikeOpConverter<mhlo::ConvolutionOp>,
+    QuantizedConvLikeOpConverter<mhlo::DynamicConvOp>
   >(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 d74b86de600..18a3fee60f0 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
@@ -1,8 +1,8 @@
 // RUN: disc-opt --disc-convert-fake-quant-op -split-input-file %s | FileCheck %s
 
-// CHECK-LABEL: @fake_quant_conv
+// CHECK-LABEL: @fake_quant_dynamic_conv
 // CHECK-SAME: (%[[INPUT:.*]]: tensor<?x32x32x6xf32>, %[[WEIGHT:.*]]: tensor<3x3x3x16xf32>
-func.func @fake_quant_conv(
+func.func @fake_quant_dynamic_conv(
     %input: tensor<?x32x32x6xf32>, %weight: tensor<3x3x3x16xf32>, %padding: tensor<4xi32>,
     %input_scale : tensor<f32>, %input_zero_point : tensor<i32>,
     %weight_scale : tensor<?xf32>, %weight_zero_point : tensor<?xi32>,
@@ -65,6 +65,77 @@ func.func @fake_quant_conv(
       use_dynamic = false
   } : (tensor<?x8x6x16xf32>, tensor<f32>, tensor<i32>) -> tensor<?x8x6x16xf32>
 
+  // CHECK: return %[[FAKE_QUANT_RESULT]] : tensor<?x8x6x16xf32>
+  return %fake_quant_result : tensor<?x8x6x16xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fake_quant_conv
+// CHECK-SAME: (%[[INPUT:.*]]: tensor<?x32x32x6xf32>, %[[WEIGHT:.*]]: tensor<3x3x3x16xf32>
+func.func @fake_quant_conv(
+    %input: tensor<?x32x32x6xf32>, %weight: tensor<3x3x3x16xf32>,
+    %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<?x8x6x16xf32> {
+  // 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<?x32x32x6xf32>, tensor<f32>, tensor<i32>) -> tensor<?x32x32x6xf32>
+
+  // 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<[2]> : tensor<1xi64>,
+      num_bits = 8,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<3x3x3x16xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<3x3x3x16xf32>
+
+  // CHECK: %[[RESULT:.*]] = "mhlo_disc.quantized_dynamic_conv"
+  // CHECK-SAME: %[[FAKE_QUANT_INPUT]]
+  // CHECK-SAME: %[[FAKE_QUANT_WEIGHT]]
+  %result = "mhlo.convolution"(%fake_quant_input, %fake_quant_weight) {
+    dimension_numbers = #mhlo.conv<raw
+      input_batch_dimension = 0,
+      input_feature_dimension = 3,
+      input_spatial_dimensions = [1, 2],
+      kernel_input_feature_dimension = 2,
+      kernel_output_feature_dimension = 3,
+      kernel_spatial_dimensions = [0, 1],
+      output_batch_dimension = 0, output_feature_dimension = 3,
+      output_spatial_dimensions = [1, 2]
+    >,
+    batch_group_count = 1 : i64,
+    feature_group_count = 2 : i64,
+    padding = dense<[[1, 2], [3, 1]]> : tensor<2x2xi64>,
+    rhs_dilation = dense<[2, 3]> : tensor<2xi64>,
+    window_strides = dense<[4, 5]> : tensor<2xi64>}
+    : (tensor<?x32x32x6xf32>, tensor<3x3x3x16xf32>) -> tensor<?x8x6x16xf32>
+
+  // 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<?x8x6x16xf32>, tensor<f32>, tensor<i32>) -> tensor<?x8x6x16xf32>
+
   // CHECK: return %[[FAKE_QUANT_RESULT]] : tensor<?x8x6x16xf32>
   return %fake_quant_result : tensor<?x8x6x16xf32>
 }
\ No newline at end of file