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/tao_bridge/ops/fake_quant_op.cc b/tao/tao_bridge/ops/fake_quant_op.cc
index 08317894a09..caf9e835ad4 100644
--- a/tao/tao_bridge/ops/fake_quant_op.cc
+++ b/tao/tao_bridge/ops/fake_quant_op.cc
@@ -24,20 +24,9 @@ class DiscFakeQuantOp : public OpKernel {
     OP_REQUIRES_OK(context, context->GetAttr("quant_max", &quant_max_));
     OP_REQUIRES_OK(context, context->GetAttr("num_bits", &num_bits_));
     OP_REQUIRES_OK(context, context->GetAttr("axis", &axis_));
-    OP_REQUIRES_OK(context, context->GetAttr("signed", &signed_));
-    OP_REQUIRES_OK(context, context->GetAttr("symmetric", &symmetric_));
-    OP_REQUIRES_OK(context, context->GetAttr("dynamic", &dynamic_));
-    OP_REQUIRES_OK(context, context->GetAttr("per_channel", &per_channel_));
-
-    if (per_channel_) {
-      OP_REQUIRES(context, axis_.size() > 0,
-                  errors::InvalidArgument(
-                      "Per-channel quantization requires non-empty axis."));
-    } else {
-      OP_REQUIRES(context, axis_.size() == 0,
-                  errors::InvalidArgument(
-                      "Per-tensor quantization requires empty axis"));
-    }
+    OP_REQUIRES_OK(context, context->GetAttr("use_signed", &use_signed_));
+    OP_REQUIRES_OK(context, context->GetAttr("use_symmetric", &use_symmetric_));
+    OP_REQUIRES_OK(context, context->GetAttr("use_dynamic", &use_dynamic_));
   }
 
   void Compute(OpKernelContext* context) override {
@@ -55,10 +44,9 @@ class DiscFakeQuantOp : public OpKernel {
   int64 quant_max_;
   int64 num_bits_;
   std::vector<int64> axis_;
-  bool signed_;
-  bool symmetric_;
-  bool dynamic_;
-  bool per_channel_;
+  bool use_signed_;
+  bool use_symmetric_;
+  bool use_dynamic_;
 };
 
 REGISTER_OP("DiscFakeQuant")
@@ -70,12 +58,11 @@ REGISTER_OP("DiscFakeQuant")
     .Attr("quant_max: int")
     .Attr("num_bits: int")
     .Attr("axis: list(int)")
-    .Attr("signed: bool")
-    .Attr("symmetric: bool")
-    .Attr("dynamic: bool")
-    .Attr("per_channel: bool")
+    .Attr("use_signed: bool")
+    .Attr("use_symmetric: bool")
+    .Attr("use_dynamic: bool")
     .Attr("Tfloat: {float}")
-    .Attr("Tint: {int64}")
+    .Attr("Tint: {int32}")
     .Doc("FakeQuant op to carry quant information. Implemented as Identity.")
     .SetShapeFn(shape_inference::UnchangedShape);
 
diff --git a/tao/tao_bridge/passes/tao_mark_for_compilation_pass.cc b/tao/tao_bridge/passes/tao_mark_for_compilation_pass.cc
index edf1d7d581e..5d4fd778b3e 100644
--- a/tao/tao_bridge/passes/tao_mark_for_compilation_pass.cc
+++ b/tao/tao_bridge/passes/tao_mark_for_compilation_pass.cc
@@ -2575,6 +2575,7 @@ std::vector<string> GetDiscSupportedOps() {
     "Const",
     "Conv2D",
     "DepthwiseConv2dNative",
+    "DiscFakeQuant",
     "DynamicStitch",
     "Equal",
     "Erf",
diff --git a/tao/tao_bridge/tf/xla_op_registry.cc b/tao/tao_bridge/tf/xla_op_registry.cc
index 150bcc31beb..665197c376a 100644
--- a/tao/tao_bridge/tf/xla_op_registry.cc
+++ b/tao/tao_bridge/tf/xla_op_registry.cc
@@ -1248,6 +1248,8 @@ REGISTER_XLA_OP_FOR_TAO(Name("QuantizedConv2DWithBiasAndRequantize"));
 #endif
 #endif
 
+REGISTER_XLA_OP_FOR_TAO(Name("DiscFakeQuant"));
+
 // clang-format on
 
 }  // namespace tao
diff --git a/tao_compiler/decoupling/BUILD b/tao_compiler/decoupling/BUILD
index a2f36a664da..e9bf6b37d52 100644
--- a/tao_compiler/decoupling/BUILD
+++ b/tao_compiler/decoupling/BUILD
@@ -71,6 +71,15 @@ tf_cc_test(
     ],
 )
 
+cc_library(
+    name = "disc_custom_ops",
+    srcs = ["fake_quant_op.cc"],
+    deps = [
+        "//tensorflow/core:framework"
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "compiler_base",
     srcs = ["compiler_base.cc"],
@@ -106,8 +115,9 @@ cc_library(
     ]),
     visibility = ["//visibility:public"],
     deps = [
-        "compiler_base",
-        "tao_compiler_trace",
+        ":compiler_base",
+        ":disc_custom_ops",
+        ":tao_compiler_trace",
         "//tensorflow/cc:cc_ops",
         "//tensorflow/cc:ops",
         "//tensorflow/cc:resource_variable_ops",
diff --git a/tao_compiler/decoupling/fake_quant_op.cc b/tao_compiler/decoupling/fake_quant_op.cc
new file mode 120000
index 00000000000..0594ed10da3
--- /dev/null
+++ b/tao_compiler/decoupling/fake_quant_op.cc
@@ -0,0 +1,1 @@
+../../tao/tao_bridge/ops/fake_quant_op.cc
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/BUILD b/tao_compiler/mlir/disc/BUILD
index 75da752ace6..b2df2ebb697 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -1168,6 +1168,7 @@ cc_library(
     ],
     deps = [
         ":disc_lower_tf_inc_gen",
+        ":disc_tf_additional_ops",
         ":disc_util",
         ":mhlo_disc",
         ":pass_details",
@@ -2334,3 +2335,43 @@ cc_library(
     ],
     alwayslink = 1,
 )
+
+gentbl_cc_library(
+    name = "disc_tf_additional_ops_gen",
+    compatible_with = get_compatible_with_cloud(),
+    strip_include_prefix = "",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "IR/disc_tf_additional_ops.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "IR/disc_tf_additional_ops.cc.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "IR/disc_tf_additional_ops.td",
+    td_srcs = [
+        "//tensorflow/compiler/mlir/tensorflow:ir/tf_op_base.td",
+        "//tensorflow/compiler/mlir/tensorflow:ir/tf_op_interfaces.td",
+    ],
+    deps = [
+        "@llvm-project//mlir:OpBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "disc_tf_additional_ops",
+    hdrs = [
+        "IR/disc_tf_additional_ops.h",
+    ],
+    srcs = [
+        "IR/disc_tf_additional_ops.cc",
+    ],
+    deps = [
+        ":disc_tf_additional_ops_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+    ],
+    alwayslink = 1,
+)
diff --git a/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.cc b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.cc
new file mode 100644
index 00000000000..13af7194d84
--- /dev/null
+++ b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.cc
@@ -0,0 +1,70 @@
+/* Copyright 2021 The TensorFlow 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 "tensorflow/compiler/mlir/disc/IR/disc_tf_additional_ops.h"
+
+namespace mlir {
+namespace TF {
+
+LogicalResult DiscFakeQuantOp::verify() {
+  auto q_min = static_cast<int64_t>(quant_min());
+  auto q_max = static_cast<int64_t>(quant_max());
+  if (q_min >= q_max) {
+    return emitOpError("quant_min (")
+           << q_min << ") must be less than quant_max (" << q_max << ")";
+  }
+
+  int64_t expect_min = use_signed() ? (1 << (num_bits() - 1)) * -1 : 0;
+  int64_t expect_max =
+      use_signed() ? (1 << (num_bits() - 1)) - 1 : (1 << num_bits()) - 1;
+
+  if (q_min < expect_min) {
+    return emitOpError("quant_min must not be less than ")
+           << expect_min << " under " << num_bits()
+           << " bits and signed=" << use_signed() << ", but got: " << q_min;
+  }
+  if (q_max > expect_max) {
+    return emitOpError("quant_max must not be greater than ")
+           << expect_max << " under " << num_bits()
+           << " bits and signed=" << use_signed() << ", but got: " << q_max;
+  }
+  if (axis().size() > 1) {
+    return emitOpError(
+               "axis must be empty (per-tensor) or a single element array "
+               "(per-channel), but got: ")
+           << axis();
+  }
+  return success();
+}
+
+namespace {
+
+void DiscTFOpsRegisterHook(TensorFlowDialect& dialect) {
+  dialect.addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/disc/IR/disc_tf_additional_ops.cc.inc"
+      >();
+}
+
+int RegisterOnce() {
+  TF_DIALECT_REGISTER_ADDITIONAL_OPERATIONS(DiscTFOpsRegisterHook);
+  return 0;
+}
+
+auto register_once = RegisterOnce();
+
+}  // namespace
+}  // namespace TF
+}  // namespace mlir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/disc/IR/disc_tf_additional_ops.cc.inc"
diff --git a/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.h b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.h
new file mode 100644
index 00000000000..158ec4edea0
--- /dev/null
+++ b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.h
@@ -0,0 +1,17 @@
+/* Copyright 2021 The TensorFlow 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 "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/disc/IR/disc_tf_additional_ops.h.inc"
diff --git a/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.td b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.td
new file mode 100644
index 00000000000..3daa1e730df
--- /dev/null
+++ b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.td
@@ -0,0 +1,61 @@
+/* Copyright 2021 The TensorFlow 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.
+==============================================================================*/
+
+
+#ifndef DISC_TF_ADDITIONAL_OPS
+#define DISC_TF_ADDITIONAL_OPS
+
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+
+include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/CallInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/IR/OpAsmInterface.td"
+include "mlir/IR/SymbolInterfaces.td"
+
+
+
+def TF_DiscFakeQuantOp : TF_Op<"DiscFakeQuant", [NoSideEffect, TF_NoConstantFold]> {
+
+  let summary = "FakeQuantOp defined by DISC to carry quantization info.";
+
+  let description = [{
+This op carry quantization info from tensorflow model to BladeDISC.
+  }];
+
+  let arguments = (ins
+    TF_Float32Tensor:$input,
+    TF_Float32Tensor:$scale,
+    TF_Int32Tensor:$zero_point,
+    DefaultValuedAttr<I64Attr, "-128">:$quant_min,
+    DefaultValuedAttr<I64Attr, "127">:$quant_max,
+    DefaultValuedAttr<I64Attr, "8">:$num_bits,
+    DefaultValuedAttr<I64ArrayAttr, "{}">:$axis,
+    DefaultValuedAttr<BoolAttr, "true">:$use_signed,
+    DefaultValuedAttr<BoolAttr, "true">:$use_symmetric,
+    DefaultValuedAttr<BoolAttr, "false">:$use_dynamic
+  );
+
+  let hasVerifier = 1;
+
+  let results = (outs
+    TF_Float32Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr Tfloat = TF_DerivedOperandTypeAttr<0>;
+  TF_DerivedOperandTypeAttr Tint = TF_DerivedOperandTypeAttr<1>;
+}
+
+#endif
diff --git a/tao_compiler/mlir/disc/transforms/lower_tf.cc b/tao_compiler/mlir/disc/transforms/lower_tf.cc
index 3c86fd1c6a6..e772c6b352f 100644
--- a/tao_compiler/mlir/disc/transforms/lower_tf.cc
+++ b/tao_compiler/mlir/disc/transforms/lower_tf.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "mlir/Transforms/DialectConversion.h"           // from @llvm-project
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
 #include "stablehlo/dialect/ChloOps.h"
+#include "tensorflow/compiler/mlir/disc/IR/disc_tf_additional_ops.h"
 #include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
 #include "tensorflow/compiler/mlir/disc/IR/rng_uniform_custom_call_op.h"
 #include "tensorflow/compiler/mlir/disc/IR/topk_custom_call_op.h"
@@ -643,6 +644,32 @@ LogicalResult ConvertDequantizeOp::dequantize(TF::DequantizeOp op,
                          << "\n";
 }
 
+struct ConvertFakeQuantOp : public OpRewritePattern<TF::DiscFakeQuantOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TF::DiscFakeQuantOp op,
+                                PatternRewriter& rewriter) const final {
+    if (op.use_dynamic()) {
+      return op->emitError("dynamic quantization is not supported yet");
+    }
+    auto loc = op.getLoc();
+    auto input = op.input();
+    auto scale = op.scale();
+    auto zero_point = op.zero_point();
+    SmallVector<int64_t> axis;
+    for (auto& v : op.axis()) {
+      axis.push_back(v.cast<IntegerAttr>().getInt());
+    }
+    Value new_op = rewriter.create<mhlo_disc::FakeQuantOp>(
+        loc, op.getType(), input, scale, zero_point, op.use_signedAttr(),
+        op.use_symmetricAttr(), GetI64ElementsAttr(axis, &rewriter),
+        op.num_bitsAttr(), op.quant_minAttr(), op.quant_maxAttr(),
+        op.use_dynamicAttr());
+    rewriter.replaceOp(op, {new_op});
+    return success();
+  }
+};
+
 Type ToLegalElementType(Type type) {
   return llvm::TypeSwitch<Type, Type>(type)
       .Case<mlir::TF::Qint8Type>([&type](Type) {
@@ -802,13 +829,6 @@ int64_t getDimSize(Type ty, int64_t index) {
   return ranked_ty.getDimSize(index);
 }
 
-mlir::DenseIntElementsAttr getI64ElementsAttr(ArrayRef<int64_t> values,
-                                              Builder* builder) {
-  RankedTensorType ty = RankedTensorType::get(
-      {static_cast<int64_t>(values.size())}, builder->getIntegerType(64));
-  return mlir::DenseIntElementsAttr::get(ty, values);
-}
-
 NamedAttribute getConvDimensionNumbersAttr(
     ArrayRef<int64_t> spatial_dims, tensorflow::TensorFormat format,
     Builder* builder,
@@ -913,10 +933,10 @@ LogicalResult parseConvParams(ConvParams& params, OpBuilder& b, Location& loc) {
   }
 
   auto rhs_dilations_attr =
-      b.getNamedAttr("rhs_dilation", getI64ElementsAttr(rhs_dilations, &b));
+      b.getNamedAttr("rhs_dilation", GetI64ElementsAttr(rhs_dilations, &b));
 
   auto window_strides_attr =
-      b.getNamedAttr("window_strides", getI64ElementsAttr(window_strides, &b));
+      b.getNamedAttr("window_strides", GetI64ElementsAttr(window_strides, &b));
 
   auto dimension_numbers_attr = getConvDimensionNumbersAttr(
       spatial_dim_indices, data_format, &b, tensorflow::FORMAT_OHWI);
@@ -1176,7 +1196,7 @@ class ConvertBucketizeOp : public OpRewritePattern<TF::BucketizeOp> {
     boradcast_dim.push_back(static_cast<int64_t>(input_rank));
     boundaries = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
         loc, broadcast_type, boundaries, broadcast_to_shape_tensor,
-        getI64ElementsAttr(boradcast_dim, &rewriter));
+        GetI64ElementsAttr(boradcast_dim, &rewriter));
 
     input = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
         loc, broadcast_type, input, broadcast_to_shape_tensor, broadcast_dims);
@@ -1193,7 +1213,7 @@ class ConvertBucketizeOp : public OpRewritePattern<TF::BucketizeOp> {
     reduce_dim.push_back(static_cast<int64_t>(input_rank));
     auto buckets = rewriter.create<mhlo::ReduceOp>(
         loc, convert.getResult(), zero,
-        getI64ElementsAttr(reduce_dim, &rewriter));
+        GetI64ElementsAttr(reduce_dim, &rewriter));
     mhlo::BuildReduceBody<mhlo::AddOp>(element_type, &buckets.body(),
                                        &rewriter);
 
@@ -1322,6 +1342,7 @@ void PrepareTFPass::runOnOperation() {
   patterns.insert<
       ConvertDequantizeOp,
       ConvertQuantizeV2Op,
+      ConvertFakeQuantOp,
       ConvertSqueezeOpDynamic,
       ConvertTopKV2OpDynamic,
       ConvertUniformOp,


Test patch:
diff --git a/tao/tao_bridge/test/gpu/mlir/test_mlir_fake_quant.py b/tao/tao_bridge/test/gpu/mlir/test_mlir_fake_quant.py
index f7e91ed56cd..91447ad02ac 100644
--- a/tao/tao_bridge/test/gpu/mlir/test_mlir_fake_quant.py
+++ b/tao/tao_bridge/test/gpu/mlir/test_mlir_fake_quant.py
@@ -28,10 +28,10 @@ class TestDiscFakeQuant(unittest.TestCase):
     def setUp(self):
         TaoTestCase._locate_lib_tao_ops()
 
-    def _run_fake_quant(self, axis, per_channel):
+    def _run_fake_quant(self, axis):
         input = np.random.rand(2,3).astype(np.float32)
         scale = np.array([0.1], dtype=np.float32)
-        zero_point = np.array([0], dtype=np.int64)
+        zero_point = np.array([0], dtype=np.int32)
         with tf.Session() as sess:
             output = TaoTestCase.LIB_TAO_OPS.disc_fake_quant(
                 tf.constant(input),
@@ -41,24 +41,17 @@ def _run_fake_quant(self, axis, per_channel):
                 quant_max=127,
                 num_bits=8,
                 axis=axis,
-                signed=True,
-                symmetric=True,
-                per_channel=per_channel,
-                dynamic=True)
+                use_signed=True,
+                use_symmetric=True,
+                use_dynamic=True)
             output = sess.run(output)
             self.assertTrue(np.allclose(output, input))
 
     def test_per_tensor(self):
-        self._run_fake_quant(axis=[], per_channel=False)
+        self._run_fake_quant(axis=[])
 
     def test_per_channel(self):
-        self._run_fake_quant(axis=[1], per_channel=True)
-
-    def test_invalid_per_tensor(self):
-        self.assertRaises(InvalidArgumentError, self._run_fake_quant, axis=[0], per_channel=False)
-
-    def test_invalid_per_channel(self):
-        self.assertRaises(InvalidArgumentError, self._run_fake_quant, axis=[], per_channel=True)
+        self._run_fake_quant(axis=[1])
 
 
 if __name__ == "__main__":
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-tf-fake-quant-invalid.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-tf-fake-quant-invalid.mlir
new file mode 100644
index 00000000000..6316b721973
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-tf-fake-quant-invalid.mlir
@@ -0,0 +1,44 @@
+// RUN: disc-opt -disc-lower-tf -split-input-file %s -verify-diagnostics
+
+
+func.func @main(%input: tensor<?x?xf32>, %scale: tensor<?xf32>, %zero_point: tensor<?xi32>) ->tensor<?x?xf32> attributes{tf.entry_function = {control_outputs = "", inputs = "arg_input_0_0_0_arg,arg_scale_0_1_0_arg,arg_zero_point_0_2_0_arg", outputs = "DiscFakeQuant:0"}} {
+  // expected-error@+1 {{dynamic quantization is not supported yet}}
+  %output = "tf.DiscFakeQuant"(%input, %scale, %zero_point) {_XlaAlreadyClustered = true, axis = [1], use_dynamic = true, use_signed = true, use_symmetric = true, num_bits = 4 : i64, quant_max = 3: i64, quant_min = -4: i64} : (tensor<?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xf32>
+  func.return %output : tensor<?x?xf32>
+}
+
+
+// -----
+
+func.func @main(%input: tensor<?x?xf32>, %scale: tensor<?xf32>, %zero_point: tensor<?xi32>) ->tensor<?x?xf32> attributes{tf.entry_function = {control_outputs = "", inputs = "arg_input_0_0_0_arg,arg_scale_0_1_0_arg,arg_zero_point_0_2_0_arg", outputs = "DiscFakeQuant:0"}} {
+  // expected-error@+1 {{quant_min (20) must be less than quant_max (10)}}
+  %output = "tf.DiscFakeQuant"(%input, %scale, %zero_point) {_XlaAlreadyClustered = true, axis = [1], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 10 : i64, quant_min = 20 : i64} : (tensor<?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xf32>
+  func.return %output : tensor<?x?xf32>
+}
+
+
+// -----
+
+func.func @main(%input: tensor<?x?xf32>, %scale: tensor<?xf32>, %zero_point: tensor<?xi32>) ->tensor<?x?xf32> attributes{tf.entry_function = {control_outputs = "", inputs = "arg_input_0_0_0_arg,arg_scale_0_1_0_arg,arg_zero_point_0_2_0_arg", outputs = "DiscFakeQuant:0"}} {
+  // expected-error@+1 {{op quant_max must not be greater than 127 under 8 bits and signed=1, but got: 130}}
+  %output = "tf.DiscFakeQuant"(%input, %scale, %zero_point) {_XlaAlreadyClustered = true, axis = [1], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 130 : i64, quant_min = -128 : i64} : (tensor<?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xf32>
+  func.return %output : tensor<?x?xf32>
+}
+
+// -----
+
+func.func @main(%input: tensor<?x?xf32>, %scale: tensor<?xf32>, %zero_point: tensor<?xi32>) ->tensor<?x?xf32> attributes{tf.entry_function = {control_outputs = "", inputs = "arg_input_0_0_0_arg,arg_scale_0_1_0_arg,arg_zero_point_0_2_0_arg", outputs = "DiscFakeQuant:0"}} {
+  // expected-error@+1 {{op quant_min must not be less than -128 under 8 bits and signed=1, but got: -130}}
+  %output = "tf.DiscFakeQuant"(%input, %scale, %zero_point) {_XlaAlreadyClustered = true, axis = [1], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 127 : i64, quant_min = -130 : i64} : (tensor<?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xf32>
+  func.return %output : tensor<?x?xf32>
+}
+
+// -----
+
+func.func @main(%input: tensor<?x?xf32>, %scale: tensor<?xf32>, %zero_point: tensor<?xi32>) ->tensor<?x?xf32> attributes{tf.entry_function = {control_outputs = "", inputs = "arg_input_0_0_0_arg,arg_scale_0_1_0_arg,arg_zero_point_0_2_0_arg", outputs = "DiscFakeQuant:0"}} {
+  // expected-error@+1 {{axis must be empty (per-tensor) or a single element array (per-channel), but got: [1, 2]}}
+  %output = "tf.DiscFakeQuant"(%input, %scale, %zero_point) {_XlaAlreadyClustered = true, axis = [1, 2], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 127 : i64, quant_min = -128 : i64} : (tensor<?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xf32>
+  func.return %output : tensor<?x?xf32>
+}
+
+
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-tf-fake-quant.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-tf-fake-quant.mlir
new file mode 100644
index 00000000000..1c041940b81
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-tf-fake-quant.mlir
@@ -0,0 +1,15 @@
+// RUN: disc-opt -disc-lower-tf -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: main
+func.func @main(%input: tensor<?x?xf32>, %scale: tensor<?xf32>, %zero_point: tensor<?xi32>) ->tensor<?x?xf32> attributes{tf.entry_function = {control_outputs = "", inputs = "arg_input_0_0_0_arg,arg_scale_0_1_0_arg,arg_zero_point_0_2_0_arg", outputs = "DiscFakeQuant:0"}} {
+  // CHECK: mhlo_disc.fake_quant
+  // CHECK-SAME: axis = dense<1>
+  // CHECK-SAME: num_bits = 4
+  // CHECK-SAME: quant_max = 3
+  // CHECK-SAME: quant_min = -4
+  // CHECK-SAME: use_dynamic = false
+  // CHECK-SAME: use_signed = true
+  // CHECK-SAME: use_symmetric = true
+  %output = "tf.DiscFakeQuant"(%input, %scale, %zero_point) {_XlaAlreadyClustered = true, axis = [1], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 4 : i64, quant_max = 3: i64, quant_min = -4: i64} : (tensor<?x?xf32>, tensor<?xf32>, tensor<?xi32>) -> tensor<?x?xf32>
+  func.return %output : tensor<?x?xf32>
+}