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/BUILD b/tao_compiler/mlir/disc/BUILD
index a5dd6463376..035c0bd92a3 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -292,11 +292,12 @@ cc_library(
     hdrs = ["transforms/placement_utils.h"],
     includes = ["include"],
     deps = [
+        ":disc_map_hlo_to_lhlo_op",
+        ":lmhlo_disc",
+        ":mhlo_disc",
         "//tensorflow/compiler/xla/mlir_hlo:mlir_hlo",
         "//tensorflow/compiler/xla/mlir_hlo:lhlo",
         "//tensorflow/compiler/xla/mlir_hlo:map_hlo_to_lhlo_op",
-        ":mhlo_disc",
-        ":lmhlo_disc",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
diff --git a/tao_compiler/mlir/disc/IR/lhlo_disc_ops.h b/tao_compiler/mlir/disc/IR/lhlo_disc_ops.h
index dfc6d4aa84b..35235c4dcfa 100644
--- a/tao_compiler/mlir/disc/IR/lhlo_disc_ops.h
+++ b/tao_compiler/mlir/disc/IR/lhlo_disc_ops.h
@@ -18,6 +18,7 @@ limitations under the License.
 #include "llvm/ADT/StringRef.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops_structs.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_structured_interface.h"
+#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
diff --git a/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td b/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td
index 3f5b02e35a6..9e6895b6662 100644
--- a/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td
+++ b/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td
@@ -30,6 +30,8 @@ def LHLODISC_Dialect : Dialect {
   let cppNamespace = "::mlir::lmhlo_disc";
 }
 
+def LHLO_I32Tensor : MemRefOf<[SignlessIntOfWidths<[32]>]>;
+
 class LHLODISC_Op<string mnemonic, list<Trait> traits> :
   Op<LHLODISC_Dialect, mnemonic,
     !listconcat([MemoryEffects<[MemRead, MemWrite]>,
@@ -93,6 +95,63 @@ def LHLODISC_PrintfOp: LHLODISC_Op<"printf", []> {
                        StrAttr:$format);
 }
 
+def LHLODISC_QuantizedDotGeneralOp: LHLODISC_Op<"quantized_dot_general", []> {
+  let summary = "quantized version of dot gerneal operator";
+  let description = [{
+    Compute the dot product using result quantized inputs.
+  }];
+  let arguments = (ins
+    Arg<LHLO_IntBuffer, "", [MemRead]>:$input,
+    Arg<LHLO_IntBuffer, "", [MemRead]>:$weight,
+    Arg<LHLO_FpBuffer, "", [MemRead]>:$input_scale,
+    Arg<LHLO_I32Tensor, "", [MemRead]>:$input_zero_point,
+    Arg<LHLO_FpBuffer, "", [MemRead]>:$weight_scale,
+    Arg<LHLO_I32Tensor, "", [MemRead]>:$weight_zero_point,
+    Arg<LHLO_FpBuffer, "", [MemRead]>:$result_scale,
+    Arg<LHLO_I32Tensor, "", [MemRead]>:$result_zero_point,
+    Arg<LHLO_IntBuffer, "", [MemWrite]>:$result,
+    DotDimensionNumbers:$dot_dimension_numbers,
+    DefaultValuedAttr<BoolAttr, "true">:$use_symmetric,
+    DefaultValuedAttr<I64ElementsAttr, "ArrayRef<int64_t>{}">:$axis,
+    DefaultValuedAttr<BoolAttr, "false">:$use_dynamic
+  );
+}
+
+def LHLODISC_QuantizedDynamicConvOp: LHLODISC_Op<"quantized_dynamic_conv", []> {
+  let summary = "quantized version of dynamic_conv operator";
+  let description = [{
+    Compute the convolution result using quantized inputs.
+  }];
+  let arguments = (ins
+    Arg<LHLO_IntBuffer, "", [MemRead]>:$input,
+    Arg<LHLO_IntBuffer, "", [MemRead]>:$weight,
+    Arg<LHLO_IntBuffer, "", [MemRead]>:$d_padding,
+    Arg<LHLO_FpBuffer, "", [MemRead]>:$input_scale,
+    Arg<LHLO_I32Tensor, "", [MemRead]>:$input_zero_point,
+    Arg<LHLO_FpBuffer, "", [MemRead]>:$weight_scale,
+    Arg<LHLO_I32Tensor, "", [MemRead]>:$weight_zero_point,
+    Arg<LHLO_FpBuffer, "", [MemRead]>:$result_scale,
+    Arg<LHLO_I32Tensor, "", [MemRead]>:$result_zero_point,
+    Arg<LHLO_IntBuffer, "", [MemWrite]>:$result,
+    // Default value: one for each of the spatial dimension.
+    OptionalAttr<I64ElementsAttr>:$window_strides,
+    // Default value: two zeros for each of the spatial dimension.
+    OptionalAttr<I64ElementsAttr>:$padding,
+    // Default value: one for each of the spatial dimension.
+    OptionalAttr<I64ElementsAttr>:$lhs_dilation,
+    // Default value: one for each of the spatial dimension.
+    OptionalAttr<I64ElementsAttr>:$rhs_dilation,
+    // Default value: false for each of the spatial dimension.
+    OptionalAttr<BoolElementsAttr>:$window_reversal,
+    ConvDimensionNumbers:$dimension_numbers,
+    I64Attr:$feature_group_count,
+    I64Attr:$batch_group_count,
+    DefaultValuedAttr<BoolAttr, "true">:$use_symmetric,
+    DefaultValuedAttr<I64ElementsAttr, "ArrayRef<int64_t>{}">:$axis,
+    DefaultValuedAttr<BoolAttr, "false">:$use_dynamic
+  );
+}
+
 def LHLODISC_SparseReshapeOp : LHLODISC_Op<"sparse_reshape", []> {
   let summary = "One-to-one mapping of TF_SparseReshapeOp";
   let description = [{
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index 7bab5bd2f17..716fee713ff 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -219,12 +219,6 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   // TODO(disc): Lower HLO shape constraints instead of eliding them here.
   pm.addNestedPass<FuncOp>(disc_ral::createDiscRemoveShapeConstraintsPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
-
-  // TODO: pay attention if this pass brings side-effects on performance
-  // TODO(disc): this pass introduces `tensor.cast(int32_tensor) ->
-  // index_tensor`, which is illegal ir. Temporarily disable this pass.
-  // TODO(disc): implement another implicit broadcast simplifier pass.
-  // pm.addNestedPass<FuncOp>(mhlo::createBroadcastPropagationPass());
   pm.addNestedPass<FuncOp>(createCSEPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
 
diff --git a/tao_compiler/mlir/disc/transforms/placement_utils.cc b/tao_compiler/mlir/disc/transforms/placement_utils.cc
index 218a0f30b3d..15a564e3e4b 100644
--- a/tao_compiler/mlir/disc/transforms/placement_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/placement_utils.cc
@@ -13,7 +13,9 @@
 
 #include "llvm/ADT/StringMap.h"
 #include "mlir-hlo/Dialect/lhlo/transforms/map_hlo_to_lhlo_op.h"
+#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
 #include "tensorflow/compiler/mlir/disc/IR/lhlo_disc_ops.h"
+#include "tensorflow/compiler/mlir/disc/transforms/disc_map_hlo_to_lhlo_op.h"
 
 namespace mlir {
 namespace placement_utils {
@@ -47,8 +49,18 @@ void appendShapeOperandListForHloOp(ShapeOperandListMap& m,
   m[TypeID::get<LhloTy>()] = list;
 }
 
+template <typename HloDiscTy>
+void appendShapeOperandListForHloDiscOp(ShapeOperandListMap& m,
+                                        ShapeOperandList list) {
+  m[TypeID::get<HloDiscTy>()] = list;
+  using LhloDiscTy = typename mhlo_disc::HloToLhloOp<HloDiscTy>;
+  m[TypeID::get<LhloDiscTy>()] = list;
+}
+
 ShapeOperandListMap initShapeCalcOperandMap() {
   ShapeOperandListMap m;
+
+  // mhlo dialect ops.
   appendShapeOperandListForHloOp<mhlo::RealDynamicSliceOp>(
       m, {/*start_indices*/ 1, /*limit_indices*/ 2, /*strides*/ 3});
   appendShapeOperandListForHloOp<mhlo::DynamicPadOp>(
@@ -62,6 +74,10 @@ ShapeOperandListMap initShapeCalcOperandMap() {
   appendShapeOperandListForHloOp<mhlo::DynamicConvOp>(m, {/*paddings*/ 2});
   appendShapeOperandListForHloOp<mhlo::IfOp>(m, {/*pred*/ 0});
 
+  // mhlo_disc dialect ops.
+  appendShapeOperandListForHloDiscOp<mhlo_disc::QuantizedDynamicConvOp>(
+      m, {/*paddings*/ 2});
+
   return m;
 }
 
diff --git a/tao_compiler/mlir/disc/transforms/placement_utils.h b/tao_compiler/mlir/disc/transforms/placement_utils.h
index fb60971a2d5..7e37e577d3f 100644
--- a/tao_compiler/mlir/disc/transforms/placement_utils.h
+++ b/tao_compiler/mlir/disc/transforms/placement_utils.h
@@ -70,11 +70,11 @@ bool isGpuMhlo(Operation* op);
 bool isGpuMemRef(Value memref);
 
 // Check Op if it is a mhlo Op.
-// TODO: remove mhlo_disc.custom_call after it is merged into mhlo dialect
 inline bool isMhloDialect(Operation* op) {
   return (op->getDialect() ==
               op->getContext()->getLoadedDialect<mhlo::MhloDialect>() ||
-          isa<mhlo_disc::CustomCallOp>(op));
+          op->getDialect() ==
+              op->getContext()->getLoadedDialect<mhlo_disc::MhloDiscDialect>());
 }
 
 inline bool isStdOnTensor(Operation* op) {


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir b/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
index 20e88e3d656..85729710f63 100644
--- a/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
@@ -1,5 +1,6 @@
 // RUN: disc-opt --disc-mhlo-mark-shape-calc -split-input-file %s | FileCheck %s
 
+// CHECK-LABEL: @main
 func.func @main(%arg0: tensor<?x8xf32>) -> tensor<?x24xf32> attributes {tf.entry_function = {input_placements = "cpu", inputs = "input0", output_placements = "cpu", outputs = "output0"}} {
   // CHECK: mhlo.constant
   // CHECK: tensor.from_elements
@@ -30,6 +31,7 @@ func.func @main(%arg0: tensor<?x8xf32>) -> tensor<?x24xf32> attributes {tf.entry
 
 // -----
 
+// CHECK-LABEL: @main
 func.func @main(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?x6x2xi64>) -> tensor<?x6x?xi32> attributes {tf.entry_function = {input_placements = "cpu,cpu", inputs = "input0, input1", output_placements = "cpu", outputs = "output0"}} {
   // CHECK: tensor.from_elements
   // CHECK-SAME: disc.shape_op = true
@@ -45,3 +47,52 @@ func.func @main(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?x6x2xi64>) -> tensor<?x
   %3 = "mhlo.dynamic_gather"(%arg0, %arg1, %2) {dimension_numbers = #mhlo.gather<collapsed_slice_dims = [0, 1], index_vector_dim = 2, offset_dims = [2], start_index_map = [0, 1]>, indices_are_sorted = false} : (tensor<?x?x?xi32>, tensor<?x6x2xi64>, tensor<3xi64>) -> tensor<?x6x?xi32>
   return %3 : tensor<?x6x?xi32>
 }
+
+// -----
+
+// Test for per_tensor_quantized_dynamic_conv
+
+// CHECK-LABEL: @main
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x?x?x?xi8>, %[[WEIGHT:.*]]: tensor<?x?x?x?xi8>,
+// CHECK-SAME: %[[PADDING:.*]]: tensor<4xf32>,
+func.func @main(%input: tensor<?x?x?x?xi8>, %weight: tensor<?x?x?x?xi8>, %padding : tensor<4xf32>,
+                %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>
+                    attributes {tf.entry_function = {
+                                  input_placements = "cpu,cpu,gpu,cpu,cpu,cpu,cpu,cpu,cpu",
+                                  inputs = "input0,input1,input2,input3,input4,input5,input6,input7,input8",
+                                  output_placements = "cpu", outputs = "output0"}} {
+  // CHECK: %[[RefinedPadding:.*]] = mhlo.convert(%[[PADDING]])
+  // CHECK-SAME: disc.shape_op = true
+  // CHECK: %[[OUT:.*]] = "mhlo_disc.quantized_dynamic_conv"
+  // CHECK-SAME: %[[INPUT]], %[[WEIGHT]], %[[RefinedPadding]]
+  %refined_padding = mhlo.convert(%padding) : (tensor<4xf32>) -> tensor<4xi32>
+  %out = "mhlo_disc.quantized_dynamic_conv"(%input, %weight, %refined_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>
+  return %out : tensor<?x?x?x?xi8>
+}
diff --git a/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir b/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir
index 85ad9f37548..c35facb0fcb 100644
--- a/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir
@@ -1,5 +1,6 @@
 // RUN: disc-opt --mhlo-place-ops -split-input-file %s | FileCheck %s
 
+// CHECK-LABEL: @main
 func.func @main(%arg : tensor<i64>) -> tensor<i64> attributes {tf.entry_function = {input_placements = "cpu", inputs = "input0", output_placements = "cpu", outputs = "output0"}}  {
   // CHECK: "mhlo.tuple"({{.*}}) {disc.device = ["cpu"]} : (tensor<i64>) -> tuple<tensor<i64>>
   // CHECK: "mhlo.get_tuple_element"({{.*}}) {disc.device = "cpu", index = 0 : i32} : (tuple<tensor<i64>>) -> tensor<i64>
@@ -10,6 +11,7 @@ func.func @main(%arg : tensor<i64>) -> tensor<i64> attributes {tf.entry_function
 
 // -----
 
+// CHECK-LABEL: @main
 func.func @main(%arg0: tensor<?x8xf32>) -> tensor<?x24xf32> attributes {tf.entry_function = {input_placements = "cpu", inputs = "input0", output_placements = "cpu", outputs = "output0"}} {
   // CHECK: mhlo.constant
   // CHECK: disc.device = "cpu"
@@ -45,6 +47,7 @@ func.func @main(%arg0: tensor<?x8xf32>) -> tensor<?x24xf32> attributes {tf.entry
 
 // -----
 
+// CHECK-LABEL: @main
 func.func @main(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?x6x2xi64>) -> tensor<?x6x?xi32> attributes {tf.entry_function = {input_placements = "cpu,cpu", inputs = "input0, input1", output_placements = "cpu", outputs = "output0"}} {
   // CHECK: tensor.from_elements
   // CHECK-NOT: disc.device = "gpu"
@@ -61,4 +64,53 @@ func.func @main(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?x6x2xi64>) -> tensor<?x
   return %3 : tensor<?x6x?xi32>
 }
 
+// -----
+
+// Test for per_tensor_quantized_dynamic_conv
+
+// CHECK-LABEL: @main
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x?x?x?xi8>, %[[WEIGHT:.*]]: tensor<?x?x?x?xi8>,
+// CHECK-SAME: %[[PADDING:.*]]: tensor<4xf32>,
+func.func @main(%input: tensor<?x?x?x?xi8>, %weight: tensor<?x?x?x?xi8>, %padding : tensor<4xf32>,
+                %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>
+                    attributes {tf.entry_function = {
+                                  input_placements = "cpu,cpu,gpu,cpu,cpu,cpu,cpu,cpu,cpu",
+                                  inputs = "input0,input1,input2,input3,input4,input5,input6,input7,input8",
+                                  output_placements = "cpu", outputs = "output0"}} {
+  // CHECK: %[[HOST_PADDING:.*]] = "mhlo_disc.d2h"(%[[PADDING]])
+  // CHECK: %[[RefinedPadding:.*]] = mhlo.convert(%[[HOST_PADDING]])
+  // CHECK-SAME: disc.shape_op = true
+  // CHECK: %[[OUT:.*]] = "mhlo_disc.quantized_dynamic_conv"
+  // CHECK-SAME: %[[RefinedPadding]]
+  %refined_padding = mhlo.convert(%padding) {disc.shape_op = true} : (tensor<4xf32>) -> tensor<4xi32>
+  %out = "mhlo_disc.quantized_dynamic_conv"(%input, %weight, %refined_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>
+  return %out : tensor<?x?x?x?xi8>
+}