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 36a1cf5841e..a5dd6463376 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -1771,6 +1771,25 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "disc_lower_quantize_and_dequantize",
+    srcs = ["transforms/disc_lower_quantize_and_dequantize.cc"],
+    deps = [
+        ":disc_shape_optimization_utils",
+        ":mhlo_disc",
+        ":pass_details",
+        "//tensorflow/compiler/xla/mlir_hlo:mlir_hlo",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "all_passes",
     hdrs = [
@@ -1800,6 +1819,7 @@ cc_library(
         ":disc_hlo_to_std",
         ":disc_lhlo_legalize_roots_to_loops",
         ":disc_lower_tf",
+        ":disc_lower_quantize_and_dequantize",
         ":disc_lower_to_library_call",
         ":disc_math_approximation",
         ":disc_memref_canonicalizer",
diff --git a/tao_compiler/mlir/disc/transforms/PassDetail.h b/tao_compiler/mlir/disc/transforms/PassDetail.h
index 4a0ee0e306e..f8d298dc9ec 100644
--- a/tao_compiler/mlir/disc/transforms/PassDetail.h
+++ b/tao_compiler/mlir/disc/transforms/PassDetail.h
@@ -46,6 +46,10 @@ namespace shape {
 class ShapeDialect;
 }
 
+namespace mhlo {
+class MhloDialect;
+}
+
 namespace disc_ral {
 
 #define GEN_PASS_CLASSES
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc b/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
new file mode 100644
index 00000000000..eac8015f8bb
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
@@ -0,0 +1,124 @@
+// Copyright 2022 The BladeDISC Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Debug.h"
+#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
+#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
+#include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
+
+#define DEBUG_TYPE "disc-lower-quantize-and-dequantize"
+
+// This file implements the logic to decompose the quantize and dequantize op
+// to a bunch of basic elementwise ops.
+
+namespace mlir {
+namespace disc_ral {
+namespace {
+
+struct QuantizeOpConverter : public OpRewritePattern<mhlo_disc::QuantizeOp> {
+  using OpRewritePattern<mhlo_disc::QuantizeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo_disc::QuantizeOp op,
+                                PatternRewriter& rewriter) const override {
+    if (op.use_dynamic())
+      return rewriter.notifyMatchFailure(op,
+                                         "Not support dynamic quantize a.t.m.");
+
+    auto inputTy = op.input().getType().cast<RankedTensorType>();
+    if (!inputTy.getElementType().isF32())
+      return rewriter.notifyMatchFailure(
+          op, "Only support quantize f32 input a.t.m.");
+
+    Location loc = op.getLoc();
+    Value inputShape = rewriter.create<shape::ShapeOfOp>(loc, op.input());
+    Value bcastedScale = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
+        loc, inputTy, op.scale(), inputShape, op.axis());
+    auto zeroPointTy = op.zero_point().getType().cast<RankedTensorType>();
+    auto castedZeroPointTy =
+        RankedTensorType::get(zeroPointTy.getShape(), inputTy.getElementType(),
+                              zeroPointTy.getEncoding());
+    Value castedZeroPoint = rewriter.create<mhlo::ConvertOp>(
+        loc, castedZeroPointTy, op.zero_point());
+    Value bcastedZeroPoint = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
+        loc, inputTy, castedZeroPoint, inputShape, op.axis());
+    auto quantMinMaxTy = RankedTensorType::get({}, inputTy.getElementType());
+    Value quantMin = rewriter.create<mhlo::ConstantOp>(
+        loc, DenseElementsAttr::get(
+                 quantMinMaxTy,
+                 static_cast<float>(static_cast<int64_t>(op.quant_min()))));
+    Value quantMax = rewriter.create<mhlo::ConstantOp>(
+        loc, DenseElementsAttr::get(
+                 quantMinMaxTy,
+                 static_cast<float>(static_cast<int64_t>(op.quant_max()))));
+    // quantMin/Max should always be scalar.
+    Value bcastedQuantMin = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
+        loc, inputTy, quantMin, inputShape, rewriter.getI64TensorAttr({}));
+    Value bcastedQuantMax = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
+        loc, inputTy, quantMax, inputShape, rewriter.getI64TensorAttr({}));
+
+    // output = clip(round(\frac{input}{scale} + zero\_point), quant\_min,
+    // quant\_max)
+    Value t0 = rewriter.create<mhlo::DivOp>(loc, op.input(), bcastedScale);
+    Value t1 = rewriter.create<mhlo::AddOp>(loc, t0, bcastedZeroPoint);
+    Value t2 = rewriter.create<mhlo::ClampOp>(loc, t1, bcastedQuantMin,
+                                              bcastedQuantMax);
+    Value out = rewriter.create<mhlo::ConvertOp>(loc, op.getType(), t2);
+    rewriter.replaceOp(op, out);
+    return success();
+  }
+};
+
+void populateQuantizeAndDequantizePatterns(RewritePatternSet& patterns) {
+  // clang-format off
+  patterns.insert<
+    QuantizeOpConverter
+  >(patterns.getContext());
+  // clang-format on
+}
+
+struct DiscLowerQuantizeAndDequantizePass
+    : public DiscLowerQuantizeAndDequantizePassBase<
+          DiscLowerQuantizeAndDequantizePass> {
+  void runOnOperation() override;
+};
+
+void DiscLowerQuantizeAndDequantizePass::runOnOperation() {
+  // Setup rewriter patterns.
+  MLIRContext& ctx = getContext();
+  RewritePatternSet patterns(&ctx);
+  populateQuantizeAndDequantizePatterns(patterns);
+
+  if (failed(
+          applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
+    signalPassFailure();
+    return;
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscLowerQuantizeAndDequantizePass() {
+  return std::make_unique<DiscLowerQuantizeAndDequantizePass>();
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/disc_passes.td b/tao_compiler/mlir/disc/transforms/disc_passes.td
index 0fdfd392a04..9fba7fef8ea 100644
--- a/tao_compiler/mlir/disc/transforms/disc_passes.td
+++ b/tao_compiler/mlir/disc/transforms/disc_passes.td
@@ -455,3 +455,13 @@ def DiscConvertFakeQuantOpPass : Pass<"disc-convert-fake-quant-op", "mlir::func:
       "disc_shape::DISCShapeDialect",
   ];
 }
+
+def DiscLowerQuantizeAndDequantizePass : Pass<"disc-lower-quantize-and-dequantize", "mlir::func::FuncOp"> {
+  let summary = "Lower quantize and dequantize op to a bound of basic elementwise ops.";
+  let constructor = "createDiscLowerQuantizeAndDequantizePass()";
+  let dependentDialects = [
+      "disc_shape::DISCShapeDialect",
+      "shape::ShapeDialect",
+      "mhlo::MhloDialect",
+  ];
+}
diff --git a/tao_compiler/mlir/disc/transforms/passes.h b/tao_compiler/mlir/disc/transforms/passes.h
index bb5c68b1a3d..db9f67487fb 100644
--- a/tao_compiler/mlir/disc/transforms/passes.h
+++ b/tao_compiler/mlir/disc/transforms/passes.h
@@ -244,6 +244,10 @@ createDiscSparseGemmTransposeSimplifierPass();
 // Converts fake_quant annotated graph to the real quantized version.
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscConvertFakeQuantOpPass();
 
+// Lowers quantize and dequantize ops to a bunch of basic elementwise ops.
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscLowerQuantizeAndDequantizePass();
+
 }  // namespace disc_ral
 }  // namespace mlir
 


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
new file mode 100644
index 00000000000..5c8d6fe65e9
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
@@ -0,0 +1,195 @@
+// RUN: disc-opt --disc-lower-quantize-and-dequantize -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @quantize_signed_per_tensor
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xf32>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<f32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<i32>
+func.func @quantize_signed_per_tensor(
+    %input: tensor<?x32x32x6xf32>,
+    %input_scale : tensor<f32>,
+    %input_zero_point : tensor<i32>) -> tensor<?x32x32x6xi8> {
+
+  // CHECK-DAG: %[[INPUT_SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK-DAG: %[[QUANT_MIN:.*]] = mhlo.constant dense<-1.280000e+02> : tensor<f32>
+  // CHECK-DAG: %[[QUANT_MAX:.*]] = mhlo.constant  dense<1.270000e+02> : tensor<f32>
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<i32>) -> tensor<f32>
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_MIN:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MIN]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_MAX:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MAX]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[T0:.*]] = mhlo.divide %[[INPUT]], %[[BCAST_SCALE]]
+  // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
+  // CHECK: %[[T2:.*]] = mhlo.clamp %[[T1]], %[[BCAST_MIN]], %[[BCAST_MAX]]
+  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T2]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xi8>
+  // CHECK: return %[[OUT]]
+
+  %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xf32>, tensor<f32>, tensor<i32>) -> tensor<?x32x32x6xi8>
+  return %quantized_input : tensor<?x32x32x6xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @quantize_unsigned_per_tensor
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xf32>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<f32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<i32>
+func.func @quantize_unsigned_per_tensor(
+    %input: tensor<?x32x32x6xf32>,
+    %input_scale : tensor<f32>,
+    %input_zero_point : tensor<i32>) -> tensor<?x32x32x6xui8> {
+
+  // CHECK-DAG: %[[INPUT_SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK-DAG: %[[QUANT_MIN:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f32>
+  // CHECK-DAG: %[[QUANT_MAX:.*]] = mhlo.constant  dense<2.550000e+02> : tensor<f32>
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<i32>) -> tensor<f32>
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_MIN:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MIN]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_MAX:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MAX]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[T0:.*]] = mhlo.divide %[[INPUT]], %[[BCAST_SCALE]]
+  // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
+  // CHECK: %[[T2:.*]] = mhlo.clamp %[[T1]], %[[BCAST_MIN]], %[[BCAST_MAX]]
+  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T2]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xui8>
+  // CHECK: return %[[OUT]]
+
+  %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[]> : tensor<0xi64>,
+      quant_min = 0,
+      quant_max = 255,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xf32>, tensor<f32>, tensor<i32>) -> tensor<?x32x32x6xui8>
+  return %quantized_input : tensor<?x32x32x6xui8>
+}
+
+// -----
+
+// CHECK-LABEL: @quantize_signed_per_channel
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xf32>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<?xf32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<?xi32>
+func.func @quantize_signed_per_channel(
+    %input: tensor<?x32x32x6xf32>,
+    %input_scale : tensor<?xf32>,
+    %input_zero_point : tensor<?xi32>) -> tensor<?x32x32x6xi8> {
+
+  // CHECK-DAG: %[[INPUT_SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK-DAG: %[[QUANT_MIN:.*]] = mhlo.constant dense<-1.280000e+02> : tensor<f32>
+  // CHECK-DAG: %[[QUANT_MAX:.*]] = mhlo.constant  dense<1.270000e+02> : tensor<f32>
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<?xi32>) -> tensor<?xf32>
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[BCAST_MIN:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MIN]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_MAX:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MAX]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[T0:.*]] = mhlo.divide %[[INPUT]], %[[BCAST_SCALE]]
+  // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
+  // CHECK: %[[T2:.*]] = mhlo.clamp %[[T1]], %[[BCAST_MIN]], %[[BCAST_MAX]]
+  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T2]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xi8>
+  // CHECK: return %[[OUT]]
+
+  %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[3]> : tensor<1xi64>,
+      quant_min = -128,
+      quant_max = 127,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x32x32x6xi8>
+  return %quantized_input : tensor<?x32x32x6xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @quantize_unsigned_per_channel
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x32x32x6xf32>
+// CHECK-SAME: %[[INPUT_SCALE:.*]]: tensor<?xf32>
+// CHECK-SAME: %[[INPUT_ZERO_POINT:.*]]: tensor<?xi32>
+func.func @quantize_unsigned_per_channel(
+    %input: tensor<?x32x32x6xf32>,
+    %input_scale : tensor<?xf32>,
+    %input_zero_point : tensor<?xi32>) -> tensor<?x32x32x6xui8> {
+
+  // CHECK-DAG: %[[INPUT_SHAPE:.*]] = shape.shape_of %[[INPUT]]
+
+  // CHECK-DAG: %[[QUANT_MIN:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f32>
+  // CHECK-DAG: %[[QUANT_MAX:.*]] = mhlo.constant  dense<2.550000e+02> : tensor<f32>
+
+  // CHECK: %[[BCAST_SCALE:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<?xi32>) -> tensor<?xf32>
+  // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
+  // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
+
+  // CHECK: %[[BCAST_MIN:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MIN]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[BCAST_MAX:.*]] = "mhlo.dynamic_broadcast_in_dim"
+  // CHECK-SAME: %[[QUANT_MAX]], %[[INPUT_SHAPE]]
+  // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
+
+  // CHECK: %[[T0:.*]] = mhlo.divide %[[INPUT]], %[[BCAST_SCALE]]
+  // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
+  // CHECK: %[[T2:.*]] = mhlo.clamp %[[T1]], %[[BCAST_MIN]], %[[BCAST_MAX]]
+  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T2]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xui8>
+  // CHECK: return %[[OUT]]
+
+  %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
+      use_symmetric = true,
+      axis = dense<[3]> : tensor<1xi64>,
+      quant_min = 0,
+      quant_max = 255,
+      use_dynamic = false
+  } : (tensor<?x32x32x6xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x32x32x6xui8>
+  return %quantized_input : tensor<?x32x32x6xui8>
+}