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/IR/hlo_disc_ops.cc b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
index a0265d265eb..2acc0c2f0e1 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
@@ -187,6 +187,42 @@ LogicalResult QuantizedDotGeneralOp::verify() {
   return CommonVerifyForQuantizedComputeIntensiveOp(this);
 }
 
+LogicalResult QuantizedDotGeneralOp::reifyReturnTypeShapes(
+    OpBuilder& builder, ValueRange operands,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
+  auto lhsType = input().getType().dyn_cast<ShapedType>();
+  auto rhsType = weight().getType().dyn_cast<ShapedType>();
+  if (!lhsType || !rhsType) {
+    return failure();
+  }
+
+  Adaptor adaptor(operands);
+  auto dimNumbers = dot_dimension_numbers();
+  SmallVector<Value> dimensions;
+  for (const int64_t lhsDim : dimNumbers.getLhsBatchingDimensions()) {
+    dimensions.push_back(
+        builder.create<tensor::DimOp>(getLoc(), adaptor.input(), lhsDim));
+  }
+
+  for (int64_t i = 0; i < lhsType.getRank(); i++) {
+    if (!llvm::is_contained(dimNumbers.getLhsContractingDimensions(), i) &&
+        !llvm::is_contained(dimNumbers.getLhsBatchingDimensions(), i)) {
+      dimensions.push_back(
+          builder.create<tensor::DimOp>(getLoc(), adaptor.input(), i));
+    }
+  }
+  for (int64_t i = 0; i < rhsType.getRank(); i++) {
+    if (!llvm::is_contained(dimNumbers.getRhsContractingDimensions(), i) &&
+        !llvm::is_contained(dimNumbers.getRhsBatchingDimensions(), i)) {
+      dimensions.push_back(
+          builder.create<tensor::DimOp>(getLoc(), adaptor.weight(), i));
+    }
+  }
+  reifiedReturnShapes.push_back(
+      builder.create<tensor::FromElementsOp>(getLoc(), dimensions));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // QuantizedDynamicConvOp
 //===----------------------------------------------------------------------===//
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
index 4be7606552e..24ced12505c 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
@@ -144,7 +144,7 @@ def HLO_DequantizeOp: HLODISC_Op<"dequantize", [NoSideEffect]> {
   let results = (outs HLO_FpTensor:$result);
 }
 
-def HLO_QuantizedDotGeneralOp: HLODISC_Op<"quantized_dot_general", [NoSideEffect]> {
+def HLO_QuantizedDotGeneralOp: HLODISC_ShapedInterfaceOp<"quantized_dot_general", [NoSideEffect]> {
   let summary = "quantized version of dot gerneal operator";
   let description = [{
     Compute the dot product using result quantized inputs.
diff --git a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
index 2f312e5cc5b..14c88297df4 100644
--- a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
@@ -250,6 +250,7 @@ void populateDiscHLOToLHLOConversionPattern(
   patterns->insert<
       HloToLhloOpConverter<mhlo_disc::H2DOp>,
       HloToLhloOpConverter<mhlo_disc::D2HOp>,
+      HloToLhloOpConverter<mhlo_disc::QuantizedDotGeneralOp>,
       HloToLhloOpConverter<mhlo_disc::QuantizedDynamicConvOp>,
       HloToLhloOpConverter<mhlo_disc::SparseReshapeOp>,
       HloToLhloOpConverter<mhlo_disc::SparseFillEmptyRowsOp>,
diff --git a/tao_compiler/mlir/disc/transforms/disc_map_hlo_to_lhlo_op.h b/tao_compiler/mlir/disc/transforms/disc_map_hlo_to_lhlo_op.h
index 246387a7c7b..abcb9a92408 100644
--- a/tao_compiler/mlir/disc/transforms/disc_map_hlo_to_lhlo_op.h
+++ b/tao_compiler/mlir/disc/transforms/disc_map_hlo_to_lhlo_op.h
@@ -39,6 +39,7 @@ using HloToLhloOp = typename HloToLhloOpImpl<HloOpTy>::Type;
 
 MAP_HLO_TO_LHLO(H2DOp);
 MAP_HLO_TO_LHLO(D2HOp);
+MAP_HLO_TO_LHLO(QuantizedDotGeneralOp);
 MAP_HLO_TO_LHLO(QuantizedDynamicConvOp);
 MAP_HLO_TO_LHLO(SparseReshapeOp);
 MAP_HLO_TO_LHLO(SparseFillEmptyRowsOp);


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-hlo-legalize-to-lhlo.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-hlo-legalize-to-lhlo.mlir
index afc2dd92a12..9ef80ceeeb7 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-hlo-legalize-to-lhlo.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-hlo-legalize-to-lhlo.mlir
@@ -112,3 +112,71 @@ func.func @per_tensor_quantized_dynamic_conv(
   // CHECK: return %[[OUT]] : memref<?x?x?x?xi8>
   return %out : tensor<?x?x?x?xi8>
 }
+
+// -----
+
+// CHECK-LABEL: @per_tensor_quantized_dot_general
+// CHECK-SAME: %[[INPUT:.*]]: memref<?x?xi8>, %[[WEIGHT:.*]]: memref<?x?xi8>
+func.func @per_tensor_quantized_dot_general(%input: tensor<?x?xi8>, %weight: tensor<?x?xi8>,
+                                            %input_scale: tensor<f32>, %input_zero_point: tensor<i32>,
+                                            %weight_scale: tensor<f32>, %weight_zero_point: tensor<i32>,
+                                            %result_scale: tensor<f32>, %result_zero_point: tensor<i32>) -> tensor<?x?xi8> {
+  // CHECK: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK: %[[M:.*]] = memref.dim %[[INPUT]], %[[C1]] : memref<?x?xi8>
+  // CHECK: %[[N:.*]] = memref.dim %[[WEIGHT]], %[[C1]] : memref<?x?xi8>
+  // CHECK: %[[RESULT:.*]] = memref.alloc(%[[M]], %[[N]]) :  memref<?x?xi8>
+  // CHECK: lmhlo_disc.quantized_dot_general
+  // CHECK-SAME: %[[INPUT]], %[[WEIGHT]]
+  // CHECK-SAME: %[[RESULT]]
+  %out = "mhlo_disc.quantized_dot_general"(%input, %weight,
+                                           %input_scale, %input_zero_point,
+                                           %weight_scale, %weight_zero_point,
+                                           %result_scale, %result_zero_point) {
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      use_dynamic = false,
+      dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>
+  } : (tensor<?x?xi8>, tensor<?x?xi8>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?xi8>
+  return %out : tensor<?x?xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @per_channel_quantized_dot_general
+// CHECK-SAME: %[[INPUT:.*]]: memref<?x?x?x?xi8>, %[[WEIGHT:.*]]: memref<?x?x?x?xi8>
+func.func @per_channel_quantized_dot_general(%input: tensor<?x?x?x?xi8>, %weight: tensor<?x?x?x?xi8>,
+                                             %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?xi8> {
+  // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+  // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+
+  // CHECK: %[[T0:.*]] = memref.dim %[[INPUT]], %[[C0]] : memref<?x?x?x?xi8>
+  // CHECK: %[[T1:.*]] = memref.dim %[[INPUT]], %[[C1]] : memref<?x?x?x?xi8>
+  // CHECK: %[[T2:.*]] = memref.dim %[[INPUT]], %[[C2]] : memref<?x?x?x?xi8>
+  // CHECK: %[[T3:.*]] = memref.dim %[[WEIGHT]], %[[C3]] : memref<?x?x?x?xi8>
+
+  // CHECK: %[[RESULT:.*]] = memref.alloc(%[[T0]], %[[T1]], %[[T2]], %[[T3]]) : memref<?x?x?x?xi8>
+  // CHECK: lmhlo_disc.quantized_dot_general
+  // CHECK-SAME: %[[INPUT]], %[[WEIGHT]]
+  // CHECK-SAME: %[[RESULT]]
+  %out = "mhlo_disc.quantized_dot_general"(%input, %weight,
+                                           %input_scale, %input_zero_point,
+                                           %weight_scale, %weight_zero_point,
+                                           %result_scale, %result_zero_point) {
+      use_symmetric = true,
+      axis = dense<[3]> : tensor<1xi64>,
+      use_dynamic = false,
+      dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0, 1], rhs_batching_dimensions = [0, 1],
+                                        lhs_contracting_dimensions = [3], rhs_contracting_dimensions = [2]>
+  } : (tensor<?x?x?x?xi8>, tensor<?x?x?x?xi8>,
+       tensor<f32>, tensor<i32>,
+       tensor<?xf32>, tensor<?xi32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?x?x?xi8>
+  return %out : tensor<?x?x?x?xi8>
+}
\ No newline at end of file