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 f02d5dc79d1..311be293176 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -433,6 +433,28 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "quantized_dot_rewriter",
+    srcs = ["transforms/quantized_dot_rewriter.cc"],
+    hdrs = ["transforms/passes.h"],
+    includes = [
+        "tensorflow/compiler/xla/mlir_hlo/include",
+        "."
+    ],
+    deps = [
+        "@com_google_absl//absl/strings",
+        ":pass_details",
+        ":placement_utils",
+        "//tensorflow/compiler/xla/mlir_hlo:mlir_hlo",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Transforms",
+        "//tensorflow/core:lib",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "mhlo_decomp_rewriters",
     srcs = ["transforms/mhlo_decomp_rewriters.cc"],
@@ -1839,6 +1861,7 @@ cc_library(
     deps = [
         ":MhloDiscPassIncGen",
         ":conv_rewriter",
+        ":quantized_dot_rewriter",
         ":disc_dot_merge",
         ":disc_algebra_simplifier",
         ":disc_assign_kernel_name",
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index c64910e7812..4a0aa818f98 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -316,6 +316,8 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(disc_ral::createDiscQuantizedConvRewriter(
       gpu_options.cc_major, gpu_options.cc_minor));
 
+  pm.addNestedPass<FuncOp>(disc_ral::createDiscQuantizedDotRewriter());
+
   if (enable_shape_constraint_ir) {
     // shape-related optimization
     pm.addPass(disc_ral::createDiscShapeOptimizationPass());
diff --git a/tao_compiler/mlir/disc/transforms/disc_passes.td b/tao_compiler/mlir/disc/transforms/disc_passes.td
index 0b4c0de2498..b150e3d4b9e 100644
--- a/tao_compiler/mlir/disc/transforms/disc_passes.td
+++ b/tao_compiler/mlir/disc/transforms/disc_passes.td
@@ -72,6 +72,11 @@ def DotRewriterPass : Pass<"disc-dot-rewriter", "mlir::func::FuncOp"> {
   let constructor = "createDiscDotRewriterPass()";
 }
 
+def QuantizedDotRewriterPass : Pass<"quantized-disc-dot-rewriter", "mlir::func::FuncOp"> {
+  let summary = "Rewrite quantized dot to transpose weight(row2col).";
+  let constructor = "createDiscQuantizedDotRewriter()";
+}
+
 def GpuConvPaddingLegalizationPass : Pass<"disc-gpu-conv-padding-legalization", "mlir::func::FuncOp"> {
   let summary = "Convert conv ops' padding value to be suitable for lowering to cudnn lib calls.";
   let constructor = "createDiscGpuConvPaddingLegalization()";
diff --git a/tao_compiler/mlir/disc/transforms/mhlo_placer.cc b/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
index aedf3217e7d..987d8f34fca 100644
--- a/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
+++ b/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
@@ -95,7 +95,8 @@ struct OpsPlacer : public PlaceOpsPassBase<OpsPlacer> {
   const DenseMap<TypeID, /*dominant operand index*/ int> kPlaceRuleMap = {
       {TypeID::get<mhlo::DynamicGatherOp>(), /*operand*/ 0},
       {TypeID::get<mhlo::GatherOp>(), /*operand*/ 0},
-      {TypeID::get<mlir::mhlo_disc::QuantizedDynamicConvOp>(), /*operand*/ 0}};
+      {TypeID::get<mlir::mhlo_disc::QuantizedDynamicConvOp>(), /*operand*/ 0},
+      {TypeID::get<mlir::mhlo_disc::QuantizedDotGeneralOp>(), /*operand*/ 0}};
 
   // Place I64 scalar output
   void placeI64ReturnedCpuScalarOps();
diff --git a/tao_compiler/mlir/disc/transforms/passes.h b/tao_compiler/mlir/disc/transforms/passes.h
index 66321f7c787..41a90cf80a2 100644
--- a/tao_compiler/mlir/disc/transforms/passes.h
+++ b/tao_compiler/mlir/disc/transforms/passes.h
@@ -60,6 +60,8 @@ std::unique_ptr<OperationPass<FuncOp>> createDiscConvRewriter(int cc_major = 8,
 std::unique_ptr<OperationPass<FuncOp>> createDiscQuantizedConvRewriter(
     int cc_major = 8, int cc_minor = 0);
 
+std::unique_ptr<OperationPass<FuncOp>> createDiscQuantizedDotRewriter();
+
 // Rewrite and decompose mhlo ops
 std::unique_ptr<OperationPass<FuncOp>>
 createDiscMhloDecompositionRewriterPass();
diff --git a/tao_compiler/mlir/disc/transforms/quantized_dot_rewriter.cc b/tao_compiler/mlir/disc/transforms/quantized_dot_rewriter.cc
new file mode 100644
index 00000000000..f855a84705d
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/quantized_dot_rewriter.cc
@@ -0,0 +1,213 @@
+// Copyright 2021 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 <unordered_set>
+
+#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"                          // TF:llvm-project
+#include "mlir/IR/Operation.h"                           // TF:llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"                      // TF:llvm-project
+#include "tensorflow/compiler/mlir/disc/disc_util.h"
+#include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
+#include "tensorflow/compiler/mlir/disc/transforms/placement_utils.h"
+
+namespace mlir {
+namespace disc_ral {
+
+namespace {
+
+// It does not support transpose of batching dimension and other dimensions
+// together now.
+// TODO: it may support the transpose between batching dimensions only, or
+// transpose of all batching dimensions together with the minor dimension.
+static inline bool isNonBatchingTransposeTensorValue(
+    Value val, SmallVector<int64_t, 4>& permutation,
+    std::unordered_set<int64_t> batching_dims) {
+  if (not val.getDefiningOp()) {
+    return false;
+  }
+  permutation.clear();
+  if (auto transpose = dyn_cast<mhlo::TransposeOp>(val.getDefiningOp())) {
+    for (auto& en :
+         llvm::enumerate(transpose.permutation().getValues<int64_t>())) {
+      if (en.index() != en.value()) {
+        if (batching_dims.find(en.index()) != batching_dims.end()) {
+          return false;
+        }
+      }
+      permutation.push_back(en.value());
+    }
+  }
+  return !permutation.empty();
+}
+
+RankedTensorType GetTransposeOutputType(
+    Value value, const SmallVectorImpl<int64_t>& transpose_permutation,
+    OpBuilder& b) {
+  // Compute the resulting shape.
+  llvm::SmallVector<int64_t, 4> transposed_shape;
+  ShapedType input_type = value.getType().cast<ShapedType>();
+  auto input_shape = input_type.getShape();
+  for (int64_t val : transpose_permutation) {
+    transposed_shape.push_back(input_shape[val]);
+  }
+  return RankedTensorType::get(transposed_shape, input_type.getElementType());
+}
+
+Operation* InsertTranspose(
+    mlir::mhlo_disc::QuantizedDotGeneralOp op, Value value,
+    const SmallVectorImpl<int64_t>& transpose_permutation, OpBuilder& b) {
+  auto transpose_permutation_attr =
+      GetI64ElementsAttr(transpose_permutation, &b);
+
+  auto transpose_type = GetTransposeOutputType(value, transpose_permutation, b);
+  auto transpose_op = b.create<mhlo::TransposeOp>(
+      op.getLoc(), transpose_type, value, transpose_permutation_attr);
+
+  if (auto attr = op->getAttr(placement_utils::kDiscPlaceAssignment))
+    transpose_op->setAttr(placement_utils::kDiscPlaceAssignment, attr);
+
+  return transpose_op;
+}
+
+// The basic logic of this pass is to insert transpose op to ensure the qgemm op
+// having expected format.
+struct QuantDotTransposeConvert
+    : public OpRewritePattern<mlir::mhlo_disc::QuantizedDotGeneralOp> {
+  explicit QuantDotTransposeConvert(MLIRContext* context)
+      : OpRewritePattern(context) {}
+
+  LogicalResult matchAndRewrite(mlir::mhlo_disc::QuantizedDotGeneralOp op,
+                                PatternRewriter& rewriter) const override {
+    bool onGpu = placement_utils::isGpuMhlo(op);
+    // Only need to transpose weight on gpu currently
+    if (onGpu == false) {
+      return failure();
+    }
+
+    auto inputTy = op.input().getType().dyn_cast<RankedTensorType>();
+    auto weightTy = op.weight().getType().dyn_cast<RankedTensorType>();
+
+    // This transpose pass can only support dot with rank equal to 2
+    if (inputTy.getRank() != 2 || weightTy.getRank() != 2) {
+      return failure();
+    }
+
+    auto dim_numbers = op.dot_dimension_numbers();
+    SmallVector<int64_t, 4> rhs_perm_check;
+    auto rhs_batching_dims_check = dim_numbers.getRhsBatchingDimensions();
+    bool tp_rhs_check = isNonBatchingTransposeTensorValue(
+        op.weight(), rhs_perm_check,
+        std::unordered_set<int64_t>(rhs_batching_dims_check.begin(),
+                                    rhs_batching_dims_check.end()));
+
+    // Match fail if weight is nxk
+    if (tp_rhs_check) {
+      return failure();
+    }
+
+    SmallVector<int64_t> transposeAttr = {1, 0};
+
+    OpBuilder b(op);
+    auto transpose_op = InsertTranspose(op, op.weight(), transposeAttr, b);
+
+    Location loc = op.getLoc();
+    Value old_lhs = op.input();
+    Value old_rhs = transpose_op->getResult(0);
+
+    RankedTensorType old_l_type =
+        old_lhs.getType().dyn_cast<RankedTensorType>();
+    RankedTensorType old_r_type =
+        old_rhs.getType().dyn_cast<RankedTensorType>();
+    if ((!old_l_type || !old_r_type)) {
+      return failure();
+    }
+
+    SmallVector<int64_t, 4> lhs_perm;
+    auto lhs_batching_dims = dim_numbers.getLhsBatchingDimensions();
+    bool tp_lhs = isNonBatchingTransposeTensorValue(
+        old_lhs, lhs_perm,
+        std::unordered_set<int64_t>(lhs_batching_dims.begin(),
+                                    lhs_batching_dims.end()));
+
+    SmallVector<int64_t, 4> rhs_perm;
+    auto rhs_batching_dims = dim_numbers.getRhsBatchingDimensions();
+    bool tp_rhs = isNonBatchingTransposeTensorValue(
+        old_rhs, rhs_perm,
+        std::unordered_set<int64_t>(rhs_batching_dims.begin(),
+                                    rhs_batching_dims.end()));
+
+    std::vector<int64_t> lhs_contracting_dims;
+    if (tp_lhs) {
+      for (auto& en :
+           llvm::enumerate(dim_numbers.getLhsContractingDimensions())) {
+        lhs_contracting_dims.push_back(lhs_perm[en.value()]);
+      }
+    } else {
+      lhs_contracting_dims = dim_numbers.getLhsContractingDimensions();
+    }
+
+    std::vector<int64_t> rhs_contracting_dims;
+    if (tp_rhs) {
+      for (auto& en :
+           llvm::enumerate(dim_numbers.getRhsContractingDimensions())) {
+        rhs_contracting_dims.push_back(rhs_perm[en.value()]);
+      }
+    } else {
+      rhs_contracting_dims = dim_numbers.getRhsContractingDimensions();
+    }
+
+    auto dot_dimension_attr = mhlo::DotDimensionNumbersAttr::get(
+        rewriter.getContext(), dim_numbers.getLhsBatchingDimensions(),
+        dim_numbers.getRhsBatchingDimensions(), lhs_contracting_dims,
+        rhs_contracting_dims);
+
+    Value newQuantizedDot =
+        rewriter.create<mlir::mhlo_disc::QuantizedDotGeneralOp>(
+            loc, op.getType(), op.input(), transpose_op->getResult(0),
+            op.input_scale(), op.input_zero_point(), op.weight_scale(),
+            op.weight_zero_point(), op.result_scale(), op.result_zero_point(),
+            dot_dimension_attr, op.use_symmetric(), op.axis(),
+            op.use_dynamic());
+
+    rewriter.replaceOp(op, newQuantizedDot);
+    return success();
+  }
+};
+
+struct DiscQuantizedDotRewriterPass
+    : public QuantizedDotRewriterPassBase<DiscQuantizedDotRewriterPass> {
+  explicit DiscQuantizedDotRewriterPass()
+      : QuantizedDotRewriterPassBase<
+            DiscQuantizedDotRewriterPass>::QuantizedDotRewriterPassBase() {}
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+
+    MLIRContext* ctx = func.getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.insert<QuantDotTransposeConvert>(ctx);
+    if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
+      func.emitError("applyPatternsAndFoldGreedily does not converge");
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createDiscQuantizedDotRewriter() {
+  return std::make_unique<DiscQuantizedDotRewriterPass>();
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/xla/ral/context/stream_executor_based_impl.cc b/tao_compiler/mlir/xla/ral/context/stream_executor_based_impl.cc
index 3e734d162b1..b24428f60ad 100644
--- a/tao_compiler/mlir/xla/ral/context/stream_executor_based_impl.cc
+++ b/tao_compiler/mlir/xla/ral/context/stream_executor_based_impl.cc
@@ -294,6 +294,9 @@ struct RalGemmState : public Context::Resource {
 #if defined(PLATFORM_ALIBABA) and defined(ENABLE_BLADE_GEMM)
 template <typename T>
 bladnn::Dtype toBlaDNNDtype() {
+  if (std::is_same<T, int8_t>::value) {
+    return bladnn::Dtype::kS8;
+  }
   if (std::is_same<T, Eigen::half>::value) {
     return bladnn::Dtype::kF16;
   }
@@ -401,6 +404,64 @@ void ral_gemm(ExecutionContext* ctx, void* stream_handle, MemRefType<InT, 2> A,
   }
 }
 
+void ral_qgemm(
+    ExecutionContext* ctx, void* stream_handle, MemRefType<int8_t, 2> input,
+    MemRefType<int8_t, 2> weight, MemRefType<float, 0> inputScales,
+    MemRefType<int32_t, 0> inputZeroPoints, MemRefType<float, 0> weightScales,
+    MemRefType<int32_t, 0> weightZeroPoints, MemRefType<float, 0> resultScales,
+    MemRefType<int32_t, 0> resultZeroPoints, MemRefType<int8_t, 2> result,
+    bool tp_a, bool tp_b, bool weight_is_const) {
+  if (isEmptyMemref(input) || isEmptyMemref(weight) || isEmptyMemref(result)) {
+    TAO_VLOG(1) << "ral_qgemm: early return for empty tensor";
+    return;
+  }
+  TAO_VLOG(0) << "tp_a: " << tp_a << " tp_b: " << tp_b;
+  TAO_VLOG(0) << "tp_a: " << tp_a << " tp_b: " << tp_b;
+  TAO_VLOG(0) << "input.sizes[0]: " << input.sizes[0]
+              << " input.sizes[1]: " << input.sizes[1];
+  TAO_VLOG(0) << "weight.sizes[0]: " << weight.sizes[0]
+              << " weight.sizes[1]: " << weight.sizes[1];
+  TAO_VLOG(0) << "result.sizes[0]: " << result.sizes[0]
+              << " result.sizes[1]: " << result.sizes[1];
+
+  auto gpu_driver = ctx->getDriver<GPUDriver>(GPUDriver::name());
+  float input_scale, weight_scale, result_scale;
+  gpu_driver->d2h(ctx, stream_handle, &inputScales.data[0], &input_scale,
+                  sizeof(float));
+  gpu_driver->d2h(ctx, stream_handle, &weightScales.data[0], &weight_scale,
+                  sizeof(float));
+  gpu_driver->d2h(ctx, stream_handle, &resultScales.data[0], &result_scale,
+                  sizeof(float));
+
+  gpu_driver->syncOnStream(ctx, stream_handle);
+
+  float kernel_scale = input_scale * weight_scale / result_scale;
+  float beta = 0.0f;
+
+#if defined(PLATFORM_ALIBABA) and defined(ENABLE_BLADE_GEMM)
+  {
+    auto gpu_driver = ctx->getDriver<GPUDriver>(GPUDriver::name());
+    auto stream =
+        static_cast<se::Stream*>(gpu_driver->asSEStream(ctx, stream_handle));
+    void* s = stream->implementation()->GpuStreamHack();
+    bladnn::Context bladnn_ctx{s};
+    bladnn::Dtype in_dtype = toBlaDNNDtype<int8_t>();
+    bladnn::Dtype out_dtype = toBlaDNNDtype<int8_t>();
+    bool ret = bladnn::gemm(
+        &bladnn_ctx, in_dtype, tp_a, input.data, input.sizes[0], input.sizes[1],
+        in_dtype, tp_b, weight.data, weight.sizes[0], weight.sizes[1],
+        out_dtype, result.data, result.sizes[0], result.sizes[1], 1, false,
+        false, &kernel_scale, &beta);
+    if (ret) {
+      return;
+    } else {
+      ctx->signalError(Context::FAILURE, "Bladnn gemm run fail.");
+    }
+  }
+#endif
+  ctx->signalError(Context::FAILURE, "Should turn on bladnn first.");
+}
+
 template <typename T, int N>
 int64_t GetBatchSize(MemRefType<T, N> memref) {
   int64_t batch = 1;
@@ -1421,7 +1482,7 @@ static bool layout_match(const std::vector<int32_t>& ref,
 
 #if defined(PLATFORM_ALIBABA) and defined(ENABLE_BLADE_GEMM)
 template <typename T, int N>
-void ral_conv_bladnn(ExecutionContext* ctx, void* stream_handle,
+bool ral_conv_bladnn(ExecutionContext* ctx, void* stream_handle,
                      MemRefType<T, N> input, MemRefType<T, N> kernel,
                      MemRefType<int32_t, 1> padding, MemRefType<T, N> output,
                      MemRefType<int32_t, 1> metadata, float alpha = 1.0) {
@@ -1716,6 +1777,7 @@ TAO_RAL_API("ral_gemm", "gpu", gpu::se_impl::ral_gemm<float, float>);
 TAO_RAL_API("ral_gemm", "gpu", gpu::se_impl::ral_gemm<double, double, double>);
 TAO_RAL_API("ral_gemm", "gpu",
             gpu::se_impl::ral_gemm<Eigen::half, Eigen::half>);
+TAO_RAL_API("ral_qgemm", "gpu", gpu::se_impl::ral_qgemm);
 TAO_RAL_API("ral_gemm", "gpu", gpu::se_impl::ral_batch_gemm<float, float, 3>);
 TAO_RAL_API("ral_gemm", "gpu", gpu::se_impl::ral_batch_gemm<float, float, 4>);
 TAO_RAL_API("ral_gemm", "gpu",


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_p_i8_per_tensor.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_p_i8_per_tensor.mlir
new file mode 100644
index 00000000000..12811b9a4ad
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_p_i8_per_tensor.mlir
@@ -0,0 +1,33 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%input: tensor<?x?xf32>) -> (tensor<?x128xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %input_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<2.0> : tensor<f32>} : () -> tensor<f32>
+      %input_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+      %quantized_input:2 = tf_executor.island wraps "tf.DiscFakeQuant"(%input, %input_scale, %input_zero_point) {
+          axis = [], 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<f32>, tensor<i32>) -> tensor<?x?xf32>
+
+      %weight:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.0> : tensor<64x128xf32>} : () -> tensor<64x128xf32>
+      %weight_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
+      %weight_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+      %quantized_weight:2 = tf_executor.island wraps "tf.DiscFakeQuant"(%weight, %weight_scale, %weight_zero_point) {
+          axis = [], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 127 : i64, quant_min = -128 : i64} :
+      (tensor<64x128xf32>, tensor<f32>, tensor<i32>) -> tensor<64x128xf32>
+
+      %result:2 = tf_executor.island wraps "tf.MatMul"(%quantized_input, %quantized_weight)
+      {
+        transpose_a = false,
+        transpose_b = false
+      } : (tensor<?x?xf32>, tensor<64x128xf32>) -> tensor<?x128xf32>
+
+      %result_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<2.0> : tensor<f32>} : () -> tensor<f32>
+      %result_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+      %quantized_result:2 = tf_executor.island wraps "tf.DiscFakeQuant"(%result, %result_scale, %result_zero_point) {
+          axis = [], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 127 : i64, quant_min = -128 : i64} :
+      (tensor<?x128xf32>, tensor<f32>, tensor<i32>) -> tensor<?x128xf32>
+
+      tf_executor.fetch %quantized_result : tensor<?x128xf32>
+    }
+    return %graph : tensor<?x128xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_matmul.cc b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_matmul.cc
index 28dfb2fb2d5..eda71e8196b 100644
--- a/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_matmul.cc
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_matmul.cc
@@ -54,4 +54,23 @@ TEST(TFQuantziedMatMul, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_CHANNEL) {
       /*expect_output_vals*/ {output}));
 }
 
+// Comment qgemm int8 GPU test since it only works when blade_gemm
+// is on
+
+// TEST(TFQuantziedMatMul, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_CHANNEL) {
+//   std::vector<float> inputs(32 * 64, 2.0);
+//   tensorflow::Tensor output(tensorflow::DataType::DT_FLOAT, {32, 128});
+//   auto datas = output.flat<float>();
+//   for (int i = 0; i < output.NumElements(); ++i) datas(i) = 128.0;
+//   EXPECT_TRUE(feature_test_main(
+//       /*mlir_file_path*/ c_ft_path + "quantized_matmul_p_i8_per_tensor.mlir",
+//       /*backend_types*/ {BackendType::kCuda},
+//       /*num_inputs*/ 1,
+//       /*num_outputs*/ 1,
+//       /*input_descriptors*/ {"32x64xf32_X"},
+//       /*output_descriptors*/ {"f32_X"},
+//       /*input_vals*/ {inputs},
+//       /*expect_output_vals*/ {output}));
+// }
+
 }  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/transforms/tests/quantized-disc-dot-rewriter.mlir b/tao_compiler/mlir/disc/transforms/tests/quantized-disc-dot-rewriter.mlir
new file mode 100644
index 00000000000..75df6b503a9
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/quantized-disc-dot-rewriter.mlir
@@ -0,0 +1,45 @@
+// RUN: disc-opt %s -quantized-disc-dot-rewriter -split-input-file | FileCheck %s
+
+// CHECK-LABEL: @main
+func.func @main(%arg0: tensor<?x?xf32>) -> tensor<?x128xf32> attributes {tf.entry_function = {input_placements = "gpu", inputs = "input0", output_placements = "gpu", outputs = "output0"}} {
+  %0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
+  %1 = mhlo.constant dense<-1.280000e+02> : tensor<f32>
+  %2 = mhlo.constant dense<1.270000e+02> : tensor<f32>
+  %3 = shape.const_shape [64, 128] : tensor<2xindex>
+  %4 = mhlo.constant dense<0.000000e+00> : tensor<128xf32>
+  %5 = mhlo.constant dense<0> : tensor<128xi32>
+  %6 = mhlo.constant dense<1.000000e+00> : tensor<128xf32>
+  %7 = mhlo.constant dense<1.000000e+00> : tensor<64x128xf32>
+  %8 = mhlo.constant dense<2.000000e+00> : tensor<f32>
+  %9 = mhlo.constant dense<0> : tensor<i32>
+  %10 = shape.shape_of %arg0 : tensor<?x?xf32> -> tensor<2xindex>
+  %11 = "mhlo.dynamic_broadcast_in_dim"(%8, %10) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<?x?xf32>
+  %12 = "mhlo.dynamic_broadcast_in_dim"(%0, %10) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<?x?xf32>
+  %13 = "mhlo.dynamic_broadcast_in_dim"(%1, %10) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<?x?xf32>
+  %14 = "mhlo.dynamic_broadcast_in_dim"(%2, %10) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<?x?xf32>
+  %15 = mhlo.divide %arg0, %11 : tensor<?x?xf32>
+  %16 = mhlo.add %15, %12 : tensor<?x?xf32>
+  %17 = mhlo.round_nearest_afz %16 : tensor<?x?xf32>
+  %18 = mhlo.clamp %17, %13, %14 : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+  %19 = mhlo.convert(%18) : (tensor<?x?xf32>) -> tensor<?x?xi8>
+  %20 = "mhlo.dynamic_broadcast_in_dim"(%6, %3) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<128xf32>, tensor<2xindex>) -> tensor<64x128xf32>
+  %21 = "mhlo.dynamic_broadcast_in_dim"(%4, %3) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<128xf32>, tensor<2xindex>) -> tensor<64x128xf32>
+  %22 = "mhlo.dynamic_broadcast_in_dim"(%1, %3) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<64x128xf32>
+  %23 = "mhlo.dynamic_broadcast_in_dim"(%2, %3) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<64x128xf32>
+  %24 = mhlo.divide %7, %20 : tensor<64x128xf32>
+  %25 = mhlo.add %24, %21 : tensor<64x128xf32>
+  %26 = mhlo.round_nearest_afz %25 : tensor<64x128xf32>
+  %27 = mhlo.clamp %26, %22, %23 : (tensor<64x128xf32>, tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+  %28 = mhlo.convert(%27) : (tensor<64x128xf32>) -> tensor<64x128xi8>
+  // CHECK: %[[T29:.*]] = "mhlo_disc.quantized_dot_general"
+  // CHECK-SAME: rhs_contracting_dimensions = [1]
+  %29 = "mhlo_disc.quantized_dot_general"(%19, %28, %8, %9, %6, %5, %8, %9) {axis = dense<1> : tensor<1xi64>, dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, use_dynamic = false, use_symmetric = true} : (tensor<?x?xi8>, tensor<64x128xi8>, tensor<f32>, tensor<i32>, tensor<128xf32>, tensor<128xi32>, tensor<f32>, tensor<i32>) -> tensor<?x128xi8>
+  %30 = shape.shape_of %29 : tensor<?x128xi8> -> tensor<2xindex>
+  %31 = "mhlo.dynamic_broadcast_in_dim"(%8, %30) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<?x128xf32>
+  %32 = "mhlo.dynamic_broadcast_in_dim"(%9, %30) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<i32>, tensor<2xindex>) -> tensor<?x128xi32>
+  %33 = mhlo.convert(%29) : (tensor<?x128xi8>) -> tensor<?x128xi32>
+  %34 = mhlo.subtract %33, %32 : tensor<?x128xi32>
+  %35 = mhlo.convert(%34) : (tensor<?x128xi32>) -> tensor<?x128xf32>
+  %36 = mhlo.multiply %35, %31 : tensor<?x128xf32>
+  return %36 : tensor<?x128xf32>
+}
\ No newline at end of file