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 d59d309da21..81a04c6b213 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
@@ -129,6 +129,71 @@ LogicalResult QuantizeOp::verify() { return QuantVerify(this); }
 
 LogicalResult DequantizeOp::verify() { return QuantVerify(this); }
 
+//===----------------------------------------------------------------------===//
+// QuantizedDotGeneralOp
+//===----------------------------------------------------------------------===//
+
+template <typename T>
+LogicalResult CommonVerifyForQuantizedComputeIntensiveOp(T* op) {
+  auto inputTy = op->input().getType().template dyn_cast<RankedTensorType>();
+  auto weightTy = op->weight().getType().template dyn_cast<RankedTensorType>();
+  auto resultTy = op->result().getType().template dyn_cast<RankedTensorType>();
+
+  if (!inputTy || !weightTy || !resultTy ||
+      inputTy.getRank() != weightTy.getRank() ||
+      inputTy.getRank() != resultTy.getRank()) {
+    return op->emitOpError()
+           << "input, weight and result should have the same rank.\n";
+  }
+
+  auto inputScaleTy =
+      op->input_scale().getType().template dyn_cast<RankedTensorType>();
+  auto inputZeroPointTy =
+      op->input_zero_point().getType().template dyn_cast<RankedTensorType>();
+  if (!inputScaleTy || !inputZeroPointTy || inputScaleTy.getRank() != 0 ||
+      inputZeroPointTy.getRank() != 0) {
+    return op->emitOpError() << "input_scale and input_zero_point only support "
+                                "per-tensor quantization\n";
+  }
+
+  auto resultScaleTy =
+      op->result_scale().getType().template dyn_cast<RankedTensorType>();
+  auto resultZeroPointTy =
+      op->result_zero_point().getType().template dyn_cast<RankedTensorType>();
+  if (!resultScaleTy || !resultZeroPointTy || resultScaleTy.getRank() != 0 ||
+      resultZeroPointTy.getRank() != 0) {
+    return op->emitOpError() << "result_scale and result_zero_point only "
+                                "support per-tensor quantization\n";
+  }
+
+  auto weightScaleTy =
+      op->weight_scale().getType().template dyn_cast<RankedTensorType>();
+  auto weightZeroPointTy =
+      op->weight_zero_point().getType().template dyn_cast<RankedTensorType>();
+  if (!weightScaleTy || !weightZeroPointTy ||
+      weightScaleTy.getShape() != weightZeroPointTy.getShape()) {
+    return op->emitOpError()
+           << "weight_scale and weight_zero_point have mismatch shape\n";
+  }
+  auto axis = op->axis().template getValues<int64_t>();
+  if (axis.size() != weightScaleTy.getRank())
+    return op->emitOpError() << "num of quantized axes (len(axis)) is not "
+                                "equal to the rank of weight_scale tensor\n";
+  return success();
+}
+
+LogicalResult QuantizedDotGeneralOp::verify() {
+  return CommonVerifyForQuantizedComputeIntensiveOp(this);
+}
+
+//===----------------------------------------------------------------------===//
+// QuantizedDynamicConvOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult QuantizedDynamicConvOp::verify() {
+  return CommonVerifyForQuantizedComputeIntensiveOp(this);
+}
+
 }  // namespace mhlo_disc
 }  // namespace mlir
 
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.h b/tao_compiler/mlir/disc/IR/hlo_disc_ops.h
index b1c8bce4af2..9cff8009bdd 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.h
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.h
@@ -16,6 +16,7 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_HLO_DISC_OPS_H_
 
 #include "llvm/ADT/StringRef.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/hlo_disc_ops.td b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
index 3b7988881f4..e9805b4f049 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
@@ -140,4 +140,61 @@ def HLO_DequantizeOp: HLODISC_Op<"dequantize", [NoSideEffect]> {
   let results = (outs HLO_FpTensor:$result);
 }
 
+def HLO_QuantizedDotGeneralOp: HLODISC_Op<"quantized_dot_general", [NoSideEffect]> {
+  let summary = "quantized version of dot gerneal operator";
+  let description = [{
+    Compute the dot product using result quantized inputs.
+  }];
+  let arguments = (ins
+    HLO_IntTensor:$input,
+    HLO_IntTensor:$weight,
+    HLO_FpTensor:$input_scale,
+    HLO_I32Tensor:$input_zero_point,
+    HLO_FpTensor:$weight_scale,
+    HLO_I32Tensor:$weight_zero_point,
+    HLO_FpTensor:$result_scale,
+    HLO_I32Tensor:$result_zero_point,
+    DotDimensionNumbers:$dot_dimension_numbers,
+    DefaultValuedAttr<BoolAttr, "true">:$use_symmetric,
+    DefaultValuedAttr<I64ElementsAttr, "ArrayRef<int64_t>{}">:$axis,
+    DefaultValuedAttr<BoolAttr, "false">:$use_dynamic
+  );
+  let results = (outs HLO_Tensor:$result);
+}
+
+def HLO_QuantizedDynamicConvOp: HLODISC_Op<"quantized_dynamic_conv", [NoSideEffect]> {
+  let summary = "quantized version of dynamic_conv operator";
+  let description = [{
+    Compute the convolution result using quantized inputs.
+  }];
+  let arguments = (ins
+    HLO_IntTensor:$input,
+    HLO_IntTensor:$weight,
+    HLO_IntTensor:$d_padding,
+    HLO_FpTensor:$input_scale,
+    HLO_I32Tensor:$input_zero_point,
+    HLO_FpTensor:$weight_scale,
+    HLO_I32Tensor:$weight_zero_point,
+    HLO_FpTensor:$result_scale,
+    HLO_I32Tensor:$result_zero_point,
+    // 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
+  );
+  let results = (outs HLO_Tensor:$result);
+}
+
 #endif // MHLO_DISC_OPS


Test patch:
diff --git a/tao_compiler/mlir/disc/IR/tests/invalid_hlo_disc.mlir b/tao_compiler/mlir/disc/IR/tests/invalid_hlo_disc.mlir
index 5c9cd4556a2..0f7969d7044 100644
--- a/tao_compiler/mlir/disc/IR/tests/invalid_hlo_disc.mlir
+++ b/tao_compiler/mlir/disc/IR/tests/invalid_hlo_disc.mlir
@@ -122,4 +122,178 @@ func.func @dequantize_mismatch_scale_and_zero_point(%input : tensor<?x?x?x?xi8>,
       use_dynamic = false
   } : (tensor<?x?x?x?xi8>, tensor<f32>, tensor<?xi32>) -> tensor<?x?x?x?xf32>
   return %out : tensor<?x?x?x?xf32>
-}
\ No newline at end of file
+}
+
+// -----
+
+func.func @mismatch_rank_quantized_dot_general(%input: tensor<?x?xi8>, %weight: tensor<?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> {
+  // expected-error@+1 {{input, weight and result should have the same rank}}
+  %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 = [1], rhs_contracting_dimensions = [0]>
+  } : (tensor<?x?xi8>, tensor<?xi8>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?xi8>
+  return %out : tensor<?x?xi8>
+}
+
+// -----
+
+func.func @quantize_per_channel_input_quantized_dot_general(%input: tensor<?x?xi8>, %weight: tensor<?x?xi8>,
+                                 %input_scale: tensor<?xf32>, %input_zero_point: tensor<?xi32>,
+                                 %weight_scale: tensor<f32>, %weight_zero_point: tensor<i32>,
+                                 %result_scale: tensor<f32>, %result_zero_point: tensor<i32>) -> tensor<?x?xi8> {
+  // expected-error@+1 {{input_scale and input_zero_point only support per-tensor quantization}}
+  %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 = [1], rhs_contracting_dimensions = [0]>
+  } : (tensor<?x?xi8>, tensor<?x?xi8>,
+       tensor<?xf32>, tensor<?xi32>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?xi8>
+  return %out : tensor<?x?xi8>
+}
+
+// -----
+
+func.func @quantize_per_channel_result_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<?xf32>, %result_zero_point: tensor<?xi32>) -> tensor<?x?xi8> {
+  // expected-error@+1 {{result_scale and result_zero_point only support per-tensor quantization}}
+  %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 = [1], rhs_contracting_dimensions = [0]>
+  } : (tensor<?x?xi8>, tensor<?x?xi8>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>,
+       tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xi8>
+  return %out : tensor<?x?xi8>
+}
+
+// -----
+
+func.func @mismatch_rank_per_tensor_quantized_dynamic_conv(%input: tensor<?x?x?x?xi8>, %weight: tensor<?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> {
+  // expected-error@+1 {{input, weight and result should have the same rank}}
+  %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?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>
+}
+
+// -----
+
+func.func @quantize_per_channel_input_quantized_dynamic_conv(%input: tensor<?x?x?x?xi8>, %weight: tensor<?x?x?x?xi8>, %padding : tensor<4xi32>,
+                                                           %input_scale: tensor<?xf32>, %input_zero_point: tensor<?xi32>,
+                                                           %weight_scale: tensor<f32>, %weight_zero_point: tensor<i32>,
+                                                           %result_scale: tensor<f32>, %result_zero_point: tensor<i32>) -> tensor<?x?x?x?xi8> {
+  // expected-error@+1 {{input_scale and input_zero_point only support per-tensor quantization}}
+  %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<?xf32>, tensor<?xi32>,
+       tensor<f32>, tensor<i32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?x?x?xi8>
+  return %out : tensor<?x?x?x?xi8>
+}
+
+// -----
+
+func.func @quantize_per_channel_result_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<?xf32>, %result_zero_point: tensor<?xi32>) -> tensor<?x?x?x?xi8> {
+  // expected-error@+1 {{result_scale and result_zero_point only support per-tensor quantization}}
+  %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<?xf32>, tensor<?xi32>) -> tensor<?x?x?x?xi8>
+  return %out : tensor<?x?x?x?xi8>
+}
diff --git a/tao_compiler/mlir/disc/IR/tests/ops_hlo_disc.mlir b/tao_compiler/mlir/disc/IR/tests/ops_hlo_disc.mlir
index 2ad5b0df483..e4619749758 100644
--- a/tao_compiler/mlir/disc/IR/tests/ops_hlo_disc.mlir
+++ b/tao_compiler/mlir/disc/IR/tests/ops_hlo_disc.mlir
@@ -81,3 +81,119 @@ func.func @dequantize_per_channel(%input : tensor<?x?x?x?xi8>, %scale : tensor<?
   } : (tensor<?x?x?x?xi8>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?x?x?xf32>
   return %out : tensor<?x?x?x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @per_tensor_quantized_dot_general
+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> {
+  %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 = [1], 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
+func.func @per_channel_quantized_dot_general(%input: tensor<?x?xi8>, %weight: tensor<?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?xi8> {
+  %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<[1]> : tensor<1xi64>,
+      use_dynamic = false,
+      dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>
+  } : (tensor<?x?xi8>, tensor<?x?xi8>,
+       tensor<f32>, tensor<i32>,
+       tensor<?xf32>, tensor<?xi32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?xi8>
+  return %out : tensor<?x?xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @per_channel_quantized_dynamic_conv
+func.func @per_channel_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<?xf32>, %weight_zero_point: tensor<?xi32>,
+                                             %result_scale: tensor<f32>, %result_zero_point: tensor<i32>) -> tensor<?x?x?x?xi8> {
+  %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<[1]> : tensor<1xi64>,
+      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<?xf32>, tensor<?xi32>,
+       tensor<f32>, tensor<i32>) -> tensor<?x?x?x?xi8>
+  return %out : tensor<?x?x?x?xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @per_tensor_quantized_dynamic_conv
+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> {
+  %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>
+  return %out : tensor<?x?x?x?xi8>
+}