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 22b43f8015b..15a991958d4 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <unordered_map>
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "tensorflow/compiler/mlir/disc/IR/custom_call_base.h"
 
@@ -194,6 +195,177 @@ LogicalResult QuantizedDynamicConvOp::verify() {
   return CommonVerifyForQuantizedComputeIntensiveOp(this);
 }
 
+//===----------------------------------------------------------------------===//
+// Mostly copy from the implementation of mhlo::DynamicConvOp
+//===----------------------------------------------------------------------===//
+
+// Returns a new scalar integer value having type `type`. Here `type` must be
+// an integer or index type.
+Value maybeCastTo(OpBuilder& b, Location loc, Value value, Type type) {
+  if (type == value.getType()) return value;
+  if (!type.isIndex() && !value.getType().isIndex()) {
+    // in case of i32 -> i64 or vice versa
+    Value casted = b.create<arith::IndexCastOp>(loc, b.getIndexType(), value);
+    return b.create<arith::IndexCastOp>(loc, type, casted);
+  }
+  return b.create<arith::IndexCastOp>(loc, type, value);
+}
+
+template <typename Op>
+LogicalResult ConvReifyReturnTypeImpl(
+    Op* op, OpBuilder& builder, ValueRange operands,
+    SmallVectorImpl<Value>& reifiedReturnShapes,
+    const SmallVector<Value>& spatial_padding_values, Type shape_scalar_type) {
+  typename Op::Adaptor adaptor(operands);
+  Value lhs = adaptor.input();
+  Value rhs = adaptor.weight();
+
+  RankedTensorType lhs_type = lhs.getType().dyn_cast<RankedTensorType>();
+  RankedTensorType rhs_type = rhs.getType().dyn_cast<RankedTensorType>();
+  // Not support unranked type a.t.m.
+  if (!lhs_type || !rhs_type) return failure();
+
+  Location loc = op->getLoc();
+
+  auto to_shape_scalar_type = [&](Value v) {
+    return maybeCastTo(builder, loc, v, shape_scalar_type);
+  };
+
+  auto dimension_numbers = op->dimension_numbers();
+  int64_t input_batch_dimension = dimension_numbers.getInputBatchDimension();
+  int64_t kernel_output_feature_dimension =
+      dimension_numbers.getKernelOutputFeatureDimension();
+  auto input_spatial_dimensions_attr =
+      dimension_numbers.getInputSpatialDimensions();
+  auto kernel_spatial_dimensions_attr =
+      dimension_numbers.getKernelSpatialDimensions();
+  auto output_spatial_dimensions_attr =
+      dimension_numbers.getOutputSpatialDimensions();
+
+  SmallVector<Value, 4> shape_values(output_spatial_dimensions_attr.size() + 2);
+
+  // batch dim = lhs-batch-dim / batch_group_count
+  Value lhs_batch_dim = to_shape_scalar_type(
+      builder.create<tensor::DimOp>(loc, lhs, input_batch_dimension));
+  Value batch_group_count = to_shape_scalar_type(
+      builder.create<arith::ConstantIndexOp>(loc, op->batch_group_count()));
+  Value batch_dim = to_shape_scalar_type(
+      builder.create<arith::DivSIOp>(loc, lhs_batch_dim, batch_group_count));
+  int64_t output_batch_dimension = dimension_numbers.getOutputBatchDimension();
+  shape_values[output_batch_dimension] = batch_dim;
+
+  // Output's feature dim is the same with kernel's output feature dim.
+  Value feature_dim = to_shape_scalar_type(
+      builder.create<tensor::DimOp>(loc, rhs, kernel_output_feature_dimension));
+  int64_t output_feature_dimension =
+      dimension_numbers.getOutputFeatureDimension();
+  shape_values[output_feature_dimension] = feature_dim;
+
+  Optional<DenseIntElementsAttr> window_strides_attr = op->window_strides();
+  Optional<DenseIntElementsAttr> lhs_dilation_attr = op->lhs_dilation();
+  Optional<DenseIntElementsAttr> rhs_dilation_attr = op->rhs_dilation();
+
+  Value one =
+      to_shape_scalar_type(builder.create<arith::ConstantIndexOp>(loc, 1));
+  for (uint64_t i = 0; i < output_spatial_dimensions_attr.size(); i++) {
+    // effective_input_value =
+    //    (input_size - 1) * input_dilation + 1 + padding_left + padding_right
+    Value effective_input_value =
+        to_shape_scalar_type(builder.create<tensor::DimOp>(
+            loc, lhs, input_spatial_dimensions_attr[i]));
+    // Dilation.
+    if (lhs_dilation_attr) {
+      Value input_dilation =
+          to_shape_scalar_type(builder.create<arith::ConstantIndexOp>(
+              loc, lhs_dilation_attr.getValue().getValues<int64_t>()[i]));
+      effective_input_value = builder.create<arith::AddIOp>(
+          loc,
+          builder.create<arith::MulIOp>(
+              loc,
+              builder.create<arith::SubIOp>(loc, effective_input_value, one),
+              input_dilation),
+          one);
+    }
+
+    // Padding.
+    if (!spatial_padding_values.empty()) {
+      Value padding_left = spatial_padding_values[i * 2];
+      Value padding_right = spatial_padding_values[i * 2 + 1];
+      effective_input_value = builder.create<arith::AddIOp>(
+          loc, effective_input_value,
+          builder.create<arith::AddIOp>(loc, padding_left, padding_right));
+    }
+
+    // effective_kernel_size = (kernel_size - 1) * dilation + 1
+    Value effective_kernel_size_value =
+        to_shape_scalar_type(builder.create<tensor::DimOp>(
+            loc, rhs, kernel_spatial_dimensions_attr[i]));
+    if (rhs_dilation_attr) {
+      Value kernel_dilation =
+          to_shape_scalar_type(builder.create<arith::ConstantIndexOp>(
+              loc, rhs_dilation_attr.getValue().getValues<int64_t>()[i]));
+      effective_kernel_size_value = builder.create<arith::AddIOp>(
+          loc, one,
+          builder.create<arith::MulIOp>(
+              loc, kernel_dilation,
+              builder.create<arith::SubIOp>(loc, effective_kernel_size_value,
+                                            one)));
+    }
+
+    // output_size =
+    //     (effective_input_value - effective_kernel_size_value) / stride + 1
+    Value output_dim_value = builder.create<arith::SubIOp>(
+        loc, effective_input_value, effective_kernel_size_value);
+    if (window_strides_attr) {
+      Value stride_value =
+          to_shape_scalar_type(builder.create<arith::ConstantIndexOp>(
+              loc, window_strides_attr.getValue().getValues<int64_t>()[i]));
+      output_dim_value =
+          builder.create<arith::DivSIOp>(loc, output_dim_value, stride_value);
+    }
+    output_dim_value =
+        builder.create<arith::AddIOp>(loc, output_dim_value, one);
+    shape_values[output_spatial_dimensions_attr[i]] = output_dim_value;
+  }
+  Value output_shape =
+      builder.create<tensor::FromElementsOp>(loc, shape_values);
+  reifiedReturnShapes.push_back(output_shape);
+  return success();
+}
+
+LogicalResult QuantizedDynamicConvOp::reifyReturnTypeShapes(
+    OpBuilder& builder, ValueRange operands,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
+  QuantizedDynamicConvOp::Adaptor adaptor(operands);
+  Value d_padding = adaptor.d_padding();
+
+  RankedTensorType padding_type =
+      d_padding.getType().dyn_cast<RankedTensorType>();
+  // Not support unranked type a.t.m.
+  if (!padding_type) return failure();
+
+  Location loc = this->getLoc();
+  Type shape_scalar_type = padding_type.getElementType();
+  auto to_shape_scalar_type = [&](Value v) {
+    return maybeCastTo(builder, loc, v, shape_scalar_type);
+  };
+
+  SmallVector<Value> spatial_padding_values;
+  auto dimension_numbers = this->dimension_numbers();
+  auto input_spatial_dimensions_attr =
+      dimension_numbers.getInputSpatialDimensions();
+  int64_t padding_num = input_spatial_dimensions_attr.size() * 2;
+  for (int64_t i = 0; i < padding_num; i++) {
+    Value offset = builder.create<arith::ConstantIndexOp>(loc, i);
+    Value pad_value = to_shape_scalar_type(
+        builder.create<tensor::ExtractOp>(loc, d_padding, offset));
+    spatial_padding_values.push_back(pad_value);
+  }
+
+  return ConvReifyReturnTypeImpl(this, builder, operands, reifiedReturnShapes,
+                                 spatial_padding_values, shape_scalar_type);
+}
+
 //===----------------------------------------------------------------------===//
 // SparseReshapeOp
 //===----------------------------------------------------------------------===//
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
index d9ce9860cb1..b6085e1b983 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
@@ -162,7 +162,7 @@ def HLO_QuantizedDotGeneralOp: HLODISC_Op<"quantized_dot_general", [NoSideEffect
   let results = (outs HLO_Tensor:$result);
 }
 
-def HLO_QuantizedDynamicConvOp: HLODISC_Op<"quantized_dynamic_conv", [NoSideEffect]> {
+def HLO_QuantizedDynamicConvOp: HLODISC_ShapedInterfaceOp<"quantized_dynamic_conv", [NoSideEffect]> {
   let summary = "quantized version of dynamic_conv operator";
   let description = [{
     Compute the convolution result using 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 9f6350ad2f0..2cf409a7508 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
@@ -237,7 +237,7 @@ struct DiscHloLegalizeToLhlo
 
   void getDependentDialects(DialectRegistry& registry) const override {
     registry.insert<lmhlo_disc::LmhloDiscDialect, memref::MemRefDialect,
-                    shape::ShapeDialect>();
+                    shape::ShapeDialect, bufferization::BufferizationDialect>();
   }
 
  public:
@@ -271,6 +271,7 @@ void populateDiscHLOToLHLOConversionPattern(
   patterns->insert<
       HloToLhloOpConverter<mhlo_disc::H2DOp>,
       HloToLhloOpConverter<mhlo_disc::D2HOp>,
+      HloToLhloOpConverter<mhlo_disc::QuantizedDynamicConvOp>,
       HloToLhloCustomCallOpConverter,
       TieShapeOpConverter,
       SparseReshapeOpConverter
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 097ae073941..793bfa0b30a 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(QuantizedDynamicConvOp);
 
 #undef MAP_HLO_TO_LHLO
 


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 95ab8291a57..afc2dd92a12 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
@@ -67,3 +67,48 @@ func.func @test_topk_custom_call(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xi64>,
   %1, %2 = "mhlo_disc.custom_call"(%arg0, %arg1, %arg2) { backend_config = "{\"dimension\": 5}", call_target_name = "topk" } : (tensor<?x?xf32>, tensor<?x?xi64>, tensor<index>) -> (tensor<?x?xf32>, tensor<?x?xi64>)
   return %1, %2 : tensor<?x?xf32>, tensor<?x?xi64>
 }
+
+// -----
+
+// CHECK-LABEL: @per_tensor_quantized_dynamic_conv
+// CHECK-SAME: %[[INPUT:.*]]: memref<?x?x?x?xi8>, %[[WEIGHT:.*]]: memref<?x?x?x?xi8>,
+// CHECK-SAME: %[[PADDING:.*]]: memref<4xi32>,
+func.func @per_tensor_quantized_dynamic_conv(
+              %input: tensor<?x?x?x?xi8>, %weight: tensor<?x?x?x?xi8>, %padding : tensor<4xi32>,
+              %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?x?x?xi8> {
+  // CHECK: %[[OUT:.*]] = memref.alloc({{.*}})
+  // CHECK-SAME: memref<?x?x?x?xi8>
+  // CHECK: "lmhlo_disc.quantized_dynamic_conv"
+  // CHECK-SAME: %[[INPUT]], %[[WEIGHT]], %[[PADDING]]
+  // CHECK-SAME: %[[OUT]]
+  %out = "mhlo_disc.quantized_dynamic_conv"(%input, %weight, %padding,
+                                           %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,
+      dimension_numbers = #mhlo.conv<raw
+        input_batch_dimension = 0,
+        input_feature_dimension = 1,
+        input_spatial_dimensions = [2, 3],
+        kernel_input_feature_dimension = 1,
+        kernel_output_feature_dimension = 0,
+        kernel_spatial_dimensions = [2, 3],
+        output_batch_dimension = 0,
+        output_feature_dimension = 1,
+        output_spatial_dimensions = [2, 3]
+      >,
+      batch_group_count = 1 : i64,
+      feature_group_count = 1 : i64,
+      rhs_dilation = dense<1> : tensor<2xi64>,
+      window_strides = dense<3> : tensor<2xi64>
+  } : (tensor<?x?x?x?xi8>, tensor<?x?x?x?xi8>, tensor<4xi32>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?x?x?xi8>
+  // CHECK: return %[[OUT]] : memref<?x?x?x?xi8>
+  return %out : tensor<?x?x?x?xi8>
+}