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/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index 25de3ca9a8b..c64910e7812 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -311,6 +311,11 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   auto& gpu_options = options.gpu_info;
   pm.addNestedPass<FuncOp>(disc_ral::createDiscConvRewriter(
       gpu_options.cc_major, gpu_options.cc_minor));
+
+  // quantize-related optimization
+  pm.addNestedPass<FuncOp>(disc_ral::createDiscQuantizedConvRewriter(
+      gpu_options.cc_major, gpu_options.cc_minor));
+
   if (enable_shape_constraint_ir) {
     // shape-related optimization
     pm.addPass(disc_ral::createDiscShapeOptimizationPass());
diff --git a/tao_compiler/mlir/disc/transforms/conv_rewriter.cc b/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
index 74ad6c06ea2..2dc2790f6e9 100644
--- a/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
+++ b/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
@@ -42,6 +42,7 @@ namespace {
 //   * in_channel, out_channel, spatial dimensions
 // - output layout: each field for one dimension. The order is:
 //   * batch, channel, spatial dimensions
+template <typename T>
 struct ConvParams {
   int num_spatial_dims;
   Value input;
@@ -53,19 +54,16 @@ struct ConvParams {
   SmallVector<int64_t> expectedInputLayout;
   SmallVector<int64_t> expectedFilterLayout;
   SmallVector<int64_t> expectedOutputLayout;
-  mhlo::DynamicConvOp conv;
+  T conv;
 };
 
-LogicalResult extractConvParams(mhlo::DynamicConvOp op, ConvParams& params) {
-  params.conv = op;
-  params.input = op.lhs();
-  params.filter = op.rhs();
-  params.output = op.getResult();
-
-  auto inputTy = params.input.getType().dyn_cast<RankedTensorType>();
-  auto filterTy = params.filter.getType().dyn_cast<RankedTensorType>();
-  auto paddingTy = op.d_padding().getType().dyn_cast<RankedTensorType>();
-  auto outputTy = params.output.getType().dyn_cast<RankedTensorType>();
+template <typename T>
+LogicalResult fillConvParams(T op, ConvParams<T>& params) {
+  auto inputTy = params.input.getType().template dyn_cast<RankedTensorType>();
+  auto filterTy = params.filter.getType().template dyn_cast<RankedTensorType>();
+  auto paddingTy =
+      op.d_padding().getType().template dyn_cast<RankedTensorType>();
+  auto outputTy = params.output.getType().template dyn_cast<RankedTensorType>();
 
   if (!inputTy || !filterTy || !paddingTy || !outputTy) {
     return op.emitOpError() << "operands must be ranked type";
@@ -103,10 +101,30 @@ LogicalResult extractConvParams(mhlo::DynamicConvOp op, ConvParams& params) {
   for (int i = 0; i < params.num_spatial_dims; ++i) {
     params.outputLayout.push_back(output_spatial_dimensions[i]);
   }
-
   return success();
 }
 
+LogicalResult extractConvParams(mhlo::DynamicConvOp op,
+                                ConvParams<mhlo::DynamicConvOp>& params) {
+  params.conv = op;
+  params.input = op.lhs();
+  params.filter = op.rhs();
+  params.output = op.getResult();
+
+  return fillConvParams<mhlo::DynamicConvOp>(op, params);
+}
+
+LogicalResult extractConvParams(
+    mlir::mhlo_disc::QuantizedDynamicConvOp op,
+    ConvParams<mlir::mhlo_disc::QuantizedDynamicConvOp>& params) {
+  params.conv = op;
+  params.input = op.input();
+  params.filter = op.weight();
+  params.output = op.getResult();
+
+  return fillConvParams<mlir::mhlo_disc::QuantizedDynamicConvOp>(op, params);
+}
+
 void fillNCHW(SmallVector<int64_t>& layout, int num_spatial_dims) {
   layout[0] = 0;
   layout[1] = 1;
@@ -147,68 +165,6 @@ void fillOHWI(SmallVector<int64_t>& layout, int num_spatial_dims) {
   }
 }
 
-// Returns true if the conv is depthwise.
-bool isDepthwiseConv(ConvParams& params) {
-  auto filterTy = params.filter.getType().cast<RankedTensorType>();
-  return filterTy.getShape()[params.filterLayout[0]] == 1;
-}
-
-LogicalResult inferExpectedLayout(ConvParams& params, int cc_major) {
-  bool onGpu = placement_utils::isGpuMhlo(params.conv);
-  auto inputTy = params.input.getType().dyn_cast<RankedTensorType>();
-  auto filterTy = params.filter.getType().dyn_cast<RankedTensorType>();
-  auto outputTy = params.output.getType().dyn_cast<RankedTensorType>();
-  int rank = inputTy.getRank();
-  int num_spatial_dims = params.num_spatial_dims;
-
-  auto& inputLayout = params.expectedInputLayout;
-  auto& filterLayout = params.expectedFilterLayout;
-  auto& outputLayout = params.expectedOutputLayout;
-  inputLayout.resize(rank);
-  filterLayout.resize(rank);
-  outputLayout.resize(rank);
-
-  if (onGpu) {
-    bool is_fp32 =
-        inputTy.getElementType().isF32() && filterTy.getElementType().isF32();
-    bool use_tf32 = true;
-    TF_CHECK_OK(tensorflow::ReadBoolFromEnvVar("NVIDIA_TF32_OVERRIDE",
-                                               /*default_val=*/use_tf32,
-                                               &use_tf32));
-    if (cc_major >= 8 && (!is_fp32 || use_tf32) ||
-        inputTy.getElementType().isF16() && filterTy.getElementType().isF16()) {
-      // TensorCore prefers NHWC layouts
-      fillNHWC(inputLayout, num_spatial_dims);
-      fillNHWC(outputLayout, num_spatial_dims);
-      fillOHWI(filterLayout, num_spatial_dims);
-    } else {
-      // Default is NCHW & OIHW
-      fillNCHW(inputLayout, num_spatial_dims);
-      fillNCHW(outputLayout, num_spatial_dims);
-      fillOIHW(filterLayout, num_spatial_dims);
-    }
-  } else {
-#if defined(TAO_AARCH64)
-    if (isDepthwiseConv(params)) {
-      fillNHWC(inputLayout, num_spatial_dims);
-      fillNHWC(outputLayout, num_spatial_dims);
-      fillHWIO(filterLayout, num_spatial_dims);
-    } else {
-      fillNHWC(inputLayout, num_spatial_dims);
-      fillNHWC(outputLayout, num_spatial_dims);
-      fillOHWI(filterLayout, num_spatial_dims);
-    }
-#else
-    // CPU conv, default layout:
-    fillNHWC(inputLayout, num_spatial_dims);
-    fillNHWC(outputLayout, num_spatial_dims);
-    fillHWIO(filterLayout, num_spatial_dims);
-#endif
-  }
-
-  return success();
-}
-
 SmallVector<int64_t> inferTransposeAttr(
     const SmallVector<int64_t>& layout,
     const SmallVector<int64_t>& expectedLayout) {
@@ -273,14 +229,80 @@ struct ConvToDynamicConvConvert : public OpRewritePattern<mhlo::ConvolutionOp> {
 
 // The basic logic of this pass is to insert transpose op to ensure the conv op
 // having expected format.
+template <typename T>
 struct DiscConvRewriterPass
-    : public ConvRewriterPassBase<DiscConvRewriterPass> {
+    : public ConvRewriterPassBase<DiscConvRewriterPass<T>> {
   explicit DiscConvRewriterPass(int cc_major, int cc_minor)
-      : ConvRewriterPassBase<DiscConvRewriterPass>::ConvRewriterPassBase() {
+      : ConvRewriterPassBase<DiscConvRewriterPass<T>>::ConvRewriterPassBase() {
     cc_major_ = cc_major;
     cc_minor_ = cc_minor;
   }
 
+  // Returns true if the conv is depthwise.
+  bool isDepthwiseConv(ConvParams<T>& params) {
+    auto filterTy = params.filter.getType().template cast<RankedTensorType>();
+    return filterTy.getShape()[params.filterLayout[0]] == 1;
+  }
+
+  LogicalResult inferExpectedLayout(ConvParams<T>& params, int cc_major) {
+    bool onGpu = placement_utils::isGpuMhlo(params.conv);
+    auto inputTy = params.input.getType().template dyn_cast<RankedTensorType>();
+    auto filterTy =
+        params.filter.getType().template dyn_cast<RankedTensorType>();
+    auto outputTy =
+        params.output.getType().template dyn_cast<RankedTensorType>();
+    int rank = inputTy.getRank();
+    int num_spatial_dims = params.num_spatial_dims;
+
+    auto& inputLayout = params.expectedInputLayout;
+    auto& filterLayout = params.expectedFilterLayout;
+    auto& outputLayout = params.expectedOutputLayout;
+    inputLayout.resize(rank);
+    filterLayout.resize(rank);
+    outputLayout.resize(rank);
+
+    if (onGpu) {
+      bool is_fp32 =
+          inputTy.getElementType().isF32() && filterTy.getElementType().isF32();
+      bool use_tf32 = true;
+      TF_CHECK_OK(tensorflow::ReadBoolFromEnvVar("NVIDIA_TF32_OVERRIDE",
+                                                 /*default_val=*/use_tf32,
+                                                 &use_tf32));
+      if (cc_major >= 8 && (!is_fp32 || use_tf32) ||
+          inputTy.getElementType().isF16() &&
+              filterTy.getElementType().isF16()) {
+        // TensorCore prefers NHWC layouts
+        fillNHWC(inputLayout, num_spatial_dims);
+        fillNHWC(outputLayout, num_spatial_dims);
+        fillOHWI(filterLayout, num_spatial_dims);
+      } else {
+        // Default is NCHW & OIHW
+        fillNCHW(inputLayout, num_spatial_dims);
+        fillNCHW(outputLayout, num_spatial_dims);
+        fillOIHW(filterLayout, num_spatial_dims);
+      }
+    } else {
+#if defined(TAO_AARCH64)
+      if (isDepthwiseConv(params)) {
+        fillNHWC(inputLayout, num_spatial_dims);
+        fillNHWC(outputLayout, num_spatial_dims);
+        fillHWIO(filterLayout, num_spatial_dims);
+      } else {
+        fillNHWC(inputLayout, num_spatial_dims);
+        fillNHWC(outputLayout, num_spatial_dims);
+        fillOHWI(filterLayout, num_spatial_dims);
+      }
+#else
+      // CPU conv, default layout:
+      fillNHWC(inputLayout, num_spatial_dims);
+      fillNHWC(outputLayout, num_spatial_dims);
+      fillHWIO(filterLayout, num_spatial_dims);
+#endif
+    }
+
+    return success();
+  }
+
   RankedTensorType GetTransposeOutputType(
       Value value, const SmallVectorImpl<int64_t>& transpose_permutation,
       OpBuilder& b) {
@@ -295,8 +317,8 @@ struct DiscConvRewriterPass
   }
 
   Operation* InsertTranspose(
-      mhlo::DynamicConvOp op, Value value,
-      const SmallVectorImpl<int64_t>& transpose_permutation, OpBuilder& b) {
+      T op, Value value, const SmallVectorImpl<int64_t>& transpose_permutation,
+      OpBuilder& b) {
     auto transpose_permutation_attr =
         GetI64ElementsAttr(transpose_permutation, &b);
 
@@ -318,7 +340,7 @@ struct DiscConvRewriterPass
     return transpose_op;
   }
 
-  void MaybeRewriteInputFormat(ConvParams& params) {
+  void MaybeRewriteInputFormat(ConvParams<T>& params) {
     if (params.inputLayout == params.expectedInputLayout) return;
 
     auto transposeAttr =
@@ -330,7 +352,7 @@ struct DiscConvRewriterPass
     params.conv.getOperation()->setOperand(0, transpose_op->getResult(0));
   }
 
-  void MaybeRewriteFilterFormat(ConvParams& params) {
+  void MaybeRewriteFilterFormat(ConvParams<T>& params) {
     if (params.filterLayout == params.expectedFilterLayout) return;
 
     auto transposeAttr =
@@ -342,7 +364,7 @@ struct DiscConvRewriterPass
     params.conv.getOperation()->setOperand(1, transpose_op->getResult(0));
   }
 
-  void MaybeRewriteOutputFormat(ConvParams& params) {
+  void MaybeRewriteOutputFormat(ConvParams<T>& params) {
     if (params.outputLayout == params.expectedOutputLayout) return;
 
     auto in2OutAttr =
@@ -363,7 +385,7 @@ struct DiscConvRewriterPass
     transpose_op->setOperand(0, params.output);
   }
 
-  void UpdateAttributes(ConvParams& params) {
+  void UpdateAttributes(ConvParams<T>& params) {
     OpBuilder b(params.conv);
     SmallVector<int64_t, 2> inputSpatialDims;
     SmallVector<int64_t, 2> filterSpatialDims;
@@ -389,8 +411,8 @@ struct DiscConvRewriterPass
             /*outputSpatialDimensions*/ outputSpatialDims));
   }
 
-  LogicalResult RewriteOp(mhlo::DynamicConvOp op) {
-    ConvParams params;
+  LogicalResult RewriteOp(T op) {
+    ConvParams<T> params;
     if (failed(extractConvParams(op, params))) {
       return failure();
     }
@@ -408,7 +430,7 @@ struct DiscConvRewriterPass
   }
 
   LogicalResult convToDynamicConv() {
-    func::FuncOp func = getOperation();
+    func::FuncOp func = this->getOperation();
     MLIRContext* ctx = func.getContext();
     RewritePatternSet patterns(ctx);
     patterns.insert<ConvToDynamicConvConvert>(ctx);
@@ -420,20 +442,22 @@ struct DiscConvRewriterPass
   }
 
   void runOnOperation() override {
-    if (failed(convToDynamicConv())) {
-      signalPassFailure();
-      return;
+    if (std::is_same<T, mhlo::DynamicConvOp>::value) {
+      if (failed(convToDynamicConv())) {
+        this->signalPassFailure();
+        return;
+      }
     }
 
-    SmallVector<mhlo::DynamicConvOp, 4> ops;
-    getOperation().walk([&](mhlo::DynamicConvOp op) { ops.push_back(op); });
+    SmallVector<T, 4> ops;
+    this->getOperation().walk([&](T op) { ops.push_back(op); });
 
     // TODO(disc): We rewrite each conv op seperately, thus may lead to
     // unnecessary transpose ops. We may implement another layout optimize pass
     // in case necessary.
     for (auto& op : ops) {
       if (failed(RewriteOp(op))) {
-        signalPassFailure();
+        this->signalPassFailure();
         return;
       }
     }
@@ -448,7 +472,15 @@ struct DiscConvRewriterPass
 
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscConvRewriter(
     int cc_major, int cc_minor) {
-  return std::make_unique<DiscConvRewriterPass>(cc_major, cc_minor);
+  return std::make_unique<DiscConvRewriterPass<mhlo::DynamicConvOp>>(cc_major,
+                                                                     cc_minor);
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>> createDiscQuantizedConvRewriter(
+    int cc_major, int cc_minor) {
+  return std::make_unique<
+      DiscConvRewriterPass<mlir::mhlo_disc::QuantizedDynamicConvOp>>(cc_major,
+                                                                     cc_minor);
 }
 
 }  // namespace disc_ral
diff --git a/tao_compiler/mlir/disc/transforms/mhlo_placer.cc b/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
index ae2457a1fb9..6cc61f5ca9f 100644
--- a/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
+++ b/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
@@ -94,7 +94,8 @@ struct OpsPlacer : public PlaceOpsPassBase<OpsPlacer> {
   // is up to the placement of the dominant operand
   const DenseMap<TypeID, /*dominant operand index*/ int> kPlaceRuleMap = {
       {TypeID::get<mhlo::DynamicGatherOp>(), /*operand*/ 0},
-      {TypeID::get<mhlo::GatherOp>(), /*operand*/ 0}};
+      {TypeID::get<mhlo::GatherOp>(), /*operand*/ 0},
+      {TypeID::get<mlir::mhlo_disc::QuantizedDynamicConvOp>(), /*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 ee286af36d9..66321f7c787 100644
--- a/tao_compiler/mlir/disc/transforms/passes.h
+++ b/tao_compiler/mlir/disc/transforms/passes.h
@@ -56,6 +56,10 @@ createDiscLhloLegalizeRootsToParallelLoopsPass(int sm_count = -1);
 std::unique_ptr<OperationPass<FuncOp>> createDiscConvRewriter(int cc_major = 8,
                                                               int cc_minor = 0);
 
+// Canonicalize qconv ops to be suitable for lowering to cudnn lib calls.
+std::unique_ptr<OperationPass<FuncOp>> createDiscQuantizedConvRewriter(
+    int cc_major = 8, int cc_minor = 0);
+
 // Rewrite and decompose mhlo ops
 std::unique_ptr<OperationPass<FuncOp>>
 createDiscMhloDecompositionRewriterPass();
diff --git a/tao_compiler/mlir/xla/ral/context/common_context_impl_cuda.cc b/tao_compiler/mlir/xla/ral/context/common_context_impl_cuda.cc
index ed65586fbff..acd0b85d9d9 100644
--- a/tao_compiler/mlir/xla/ral/context/common_context_impl_cuda.cc
+++ b/tao_compiler/mlir/xla/ral/context/common_context_impl_cuda.cc
@@ -472,6 +472,15 @@ RAL_REGISTER_GPU_COPY_MEMREF_FUNC(Eigen::half, 5);
 RAL_REGISTER_GPU_COPY_MEMREF_FUNC(Eigen::half, 6);
 RAL_REGISTER_GPU_COPY_MEMREF_FUNC(Eigen::half, 7);
 RAL_REGISTER_GPU_COPY_MEMREF_FUNC(Eigen::half, 8);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC_0D(int8_t);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 1);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 2);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 3);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 4);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 5);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 6);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 7);
+RAL_REGISTER_GPU_COPY_MEMREF_FUNC(int8_t, 8);
 #endif  // TAO_RAL_USE_STREAM_EXECUTOR
 
 }  // namespace ral
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 842b56d98d6..3e734d162b1 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
@@ -588,11 +588,13 @@ struct CudnnConvParams {
   int64_t algo_id;
   size_t best_result_bytes_used;
   bool tensor_ops_enabled;
+  float scale;
 
   BatchDescriptor input_descriptor;
   FilterDescriptor filter_descriptor;
   ConvolutionDescriptor convolution_descriptor;
   BatchDescriptor output_descriptor;
+  BatchDescriptor bias_descriptor;
 };
 
 #if TENSORFLOW_USE_ROCM
@@ -1257,6 +1259,89 @@ Status RunCudnnConvolution(CudnnConvParams& params,
   return Status::OK();
 }
 
+// partial specialization for int8 kernel
+template <>
+Status RunCudnnConvolution<int8_t>(
+    CudnnConvParams& params, std::vector<se::DeviceMemoryBase>& operand_buffers,
+    se::DeviceMemoryBase& result_buffer,
+    se::ScratchAllocator* scratch_allocator, se::Stream* stream,
+    se::dnn::ProfileResult* profile_result) {
+  ConvolutionKind kind = params.kind;
+  DeviceMemory<int8_t> input_buf(operand_buffers[0]);
+  DeviceMemory<int8_t> filter_buf(operand_buffers[1]);
+  DeviceMemory<int8_t> output_buf(result_buffer);
+  auto& input_descriptor = params.input_descriptor;
+  auto& filter_descriptor = params.filter_descriptor;
+  auto& convolution_descriptor = params.convolution_descriptor;
+  auto& output_descriptor = params.output_descriptor;
+  AlgorithmConfig algorithm{
+      AlgorithmDesc(params.algo_id, params.tensor_ops_enabled)};
+#if TENSORFLOW_USE_ROCM
+  if (profile_result) {
+    algorithm.set_scratch_size(profile_result->scratch_size());
+  } else {
+    algorithm.set_scratch_size(params.best_result_bytes_used);
+  }
+#endif
+
+  if (TAO_VLOG_IS_ON(2)) {
+    TAO_VLOG(0) << "input ptr: " << input_buf.opaque() << "@"
+                << input_buf.size();
+    TAO_VLOG(0) << "filter_buf ptr: " << filter_buf.opaque() << "@"
+                << filter_buf.size();
+    TAO_VLOG(0) << "output_buf ptr: " << output_buf.opaque() << "@"
+                << output_buf.size();
+    TAO_VLOG(0) << "\tdescriptors:\n"
+                << "\t  input: " << params.input_descriptor.ToString() << "\n"
+                << "\t  filter: " << params.filter_descriptor.ToString() << "\n"
+                << "\t  conv: " << params.convolution_descriptor.ToString()
+                << "\n"
+                << "\t  output: " << params.output_descriptor.ToString()
+                << "\n";
+    TAO_VLOG(0) << "\talgorithm: " << algorithm.ToString();
+    TAO_VLOG(0) << "\tprofile_result: " << profile_result;
+    TAO_VLOG(0) << "\tscratch: " << scratch_allocator;
+  }
+
+  Status status = Status::OK();
+
+  switch (kind) {
+    case ConvolutionKind::FORWARD:
+
+#if (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION > 6) || TF_MAJOR_VERSION > 2
+      stream->FusedConvolveWithAlgorithm(
+          input_descriptor, input_buf, params.scale, filter_descriptor,
+          filter_buf, convolution_descriptor, output_buf, (float)0.0,
+          params.bias_descriptor, DeviceMemory<int32_t>(operand_buffers[2]),
+          se::dnn::ActivationMode::kNone, output_descriptor, &output_buf,
+          scratch_allocator, algorithm, profile_result);
+#elif TF_MAJOR_VERSION > 1
+      stream->FusedConvolveWithAlgorithm(
+          input_descriptor, input_buf, params.scale, filter_descriptor,
+          filter_buf, convolution_descriptor, output_buf, (float)0.0,
+          params.bias_descriptor, DeviceMemory<float>(operand_buffers[2]),
+          se::dnn::ActivationMode::kNone, output_descriptor, &output_buf,
+          scratch_allocator, algorithm, profile_result);
+#else
+      stream->ThenFusedConvolveWithAlgorithm(
+          input_descriptor, input_buf, params.scale, filter_descriptor,
+          filter_buf, convolution_descriptor, output_buf, (float)0.0,
+          params.bias_descriptor, DeviceMemory<int32_t>(operand_buffers[2]),
+          se::dnn::ActivationMode::kNone, output_descriptor, &output_buf,
+          scratch_allocator, algorithm, profile_result);
+#endif
+      break;
+    default:
+      return errors::Internal("Not known CudnnConvKind");
+  }
+
+  if (!status.ok()) return status;
+  if (!stream->ok()) {
+    return errors::Internal("Unable to launch convolution");
+  }
+  return Status::OK();
+}
+
 template <typename T>
 bool PickBestAlgorithm(CudnnConvParams& params,
                        std::vector<se::DeviceMemoryBase>& operand_buffers,
@@ -1334,34 +1419,12 @@ static bool layout_match(const std::vector<int32_t>& ref,
   return true;
 }
 
+#if defined(PLATFORM_ALIBABA) and defined(ENABLE_BLADE_GEMM)
 template <typename T, int N>
-void ral_conv(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) {
-  static_assert(N > 2, "dimension should be larger than 2");
-  if (isEmptyMemref(input) || isEmptyMemref(kernel) || isEmptyMemref(output)) {
-    TAO_VLOG(1) << "ral_conv: early return for empty tensor";
-    return;
-  }
-
-  if (TAO_VLOG_IS_ON(2)) {
-    print_memref(input, "input");
-    print_memref(kernel, "kernel");
-    print_memref(padding, "padding");
-    for (int i = 0; i < N - 2; ++i) {
-      TAO_VLOG(0) << "\tpadding for dim #" << i << ": (" << padding.data[2 * i]
-                  << ", " << padding.data[2 * i + 1] << ")";
-    }
-    print_memref(output, "output");
-    print_memref(metadata, "metadata");
-    for (int i = 0; i < 5 * N - 4; ++i) {
-      TAO_VLOG(0) << "\t#" << i << ": " << metadata.data[i];
-    }
-  }
-
-#if defined(PLATFORM_ALIBABA) and defined(ENABLE_BLADE_GEMM) and \
-    !(TENSORFLOW_USE_ROCM)
+void 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) {
   // nchw=0123 iohw=0123
   // const std::vector<int32_t> nchw_oihw_layout = {0, 1, 2, 3, 1, 0,
   //                                                2, 3, 0, 1, 2, 3};
@@ -1417,7 +1480,6 @@ void ral_conv(ExecutionContext* ctx, void* stream_handle,
     auto conv_kind = bladnn::ConvKind::kFprop;
     auto data_layout = bladnn::Layout::kNHWC;
     auto kernel_layout = bladnn::Layout::kNHWC;
-    const float alpha = 1.0f;
     const float beta = 0.0f;
     bladnn::Dtype dtype = toBlaDNNDtype<T>();
     bool ret = false;
@@ -1425,10 +1487,44 @@ void ral_conv(ExecutionContext* ctx, void* stream_handle,
                          n, ih, iw, ic, ko, kh, kw, oh, ow, pad_h, pad_w,
                          stride_h, stride_w, dilation_h, dilation_w, groups,
                          &alpha, a_data, b_data, &beta, c_data, c_data);
-    if (ret) {
-      return;
+    return ret;
+  }
+  return false;
+}
+#endif
+
+template <typename T, int N>
+void ral_conv(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) {
+  static_assert(N > 2, "dimension should be larger than 2");
+  if (isEmptyMemref(input) || isEmptyMemref(kernel) || isEmptyMemref(output)) {
+    TAO_VLOG(1) << "ral_conv: early return for empty tensor";
+    return;
+  }
+
+  if (TAO_VLOG_IS_ON(2)) {
+    print_memref(input, "input");
+    print_memref(kernel, "kernel");
+    print_memref(padding, "padding");
+    for (int i = 0; i < N - 2; ++i) {
+      TAO_VLOG(0) << "\tpadding for dim #" << i << ": (" << padding.data[2 * i]
+                  << ", " << padding.data[2 * i + 1] << ")";
+    }
+    print_memref(output, "output");
+    print_memref(metadata, "metadata");
+    for (int i = 0; i < 5 * N - 4; ++i) {
+      TAO_VLOG(0) << "\t#" << i << ": " << metadata.data[i];
     }
   }
+
+#if defined(PLATFORM_ALIBABA) and defined(ENABLE_BLADE_GEMM) and \
+    !(TENSORFLOW_USE_ROCM)
+  if (ral_conv_bladnn<T, N>(ctx, stream_handle, input, kernel, padding, output,
+                            metadata)) {
+    return;
+  }
 #endif
 
   auto key = makeConvTuningCacheKey(input, kernel, padding, output, metadata);
@@ -1485,6 +1581,126 @@ void ral_conv(ExecutionContext* ctx, void* stream_handle,
   return;
 }
 
+template <typename T, int N>
+void ral_qconv(ExecutionContext* ctx, void* stream_handle,
+               MemRefType<T, N> input, MemRefType<T, N> kernel,
+               MemRefType<int32_t, 1> padding, MemRefType<float, 0> inputScales,
+               MemRefType<int32_t, 0> inputZeroPoints,
+               MemRefType<float, 1> weightScales,
+               MemRefType<int32_t, 1> weightZeroPoints,
+               MemRefType<float, 0> resultScales,
+               MemRefType<int32_t, 0> resultZeroPoints, MemRefType<T, N> output,
+               MemRefType<int32_t, 1> metadata) {
+  static_assert(N > 2, "dimension should be larger than 2");
+  if (isEmptyMemref(input) || isEmptyMemref(kernel) || isEmptyMemref(output)) {
+    TAO_VLOG(1) << "ral_conv: early return for empty tensor";
+    return;
+  }
+
+  if (TAO_VLOG_IS_ON(2)) {
+    print_memref(input, "input");
+    print_memref(kernel, "kernel");
+    print_memref(padding, "padding");
+    for (int i = 0; i < N - 2; ++i) {
+      TAO_VLOG(0) << "\tpadding for dim #" << i << ": (" << padding.data[2 * i]
+                  << ", " << padding.data[2 * i + 1] << ")";
+    }
+    print_memref(output, "output");
+    print_memref(metadata, "metadata");
+    for (int i = 0; i < 5 * N - 4; ++i) {
+      TAO_VLOG(0) << "\t#" << i << ": " << metadata.data[i];
+    }
+  }
+
+  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;
+
+#if defined(PLATFORM_ALIBABA) and defined(ENABLE_BLADE_GEMM) and \
+    !(TENSORFLOW_USE_ROCM)
+  if (ral_conv_bladnn<T, N>(ctx, stream_handle, input, kernel, padding, output,
+                            metadata, kernel_scale)) {
+    return;
+  }
+#endif
+
+  auto key = makeConvTuningCacheKey(input, kernel, padding, output, metadata);
+  CudnnConvParams* params = nullptr;
+  std::string unique_name =
+      "tao_ral.gpu.qconv_" + tao::ral::TaoTypeNameHelper<T>::Invoke();
+  auto& state = RalConvState::Get();
+  auto stream =
+      static_cast<se::Stream*>(gpu_driver->asSEStream(ctx, stream_handle));
+
+  std::vector<se::DeviceMemoryBase> operand_se_buffers;
+  operand_se_buffers.emplace_back(GetDeviceAddress(input));
+  operand_se_buffers.emplace_back(GetDeviceAddress(kernel));
+  // buffer for bias
+  operand_se_buffers.emplace_back(GetDeviceAddress(weightZeroPoints));
+  se::DeviceMemoryBase result_buffer = GetDeviceAddress(output);
+
+  {
+    std::lock_guard<std::mutex> l(state.mu);
+    auto& cache = state.cache_table[unique_name];
+    auto it = cache.find(key);
+    if (it == cache.end()) {
+      auto params_ptr = makeNewConvParams(key);
+      auto& params = *params_ptr;
+      params.scale = kernel_scale;
+      BatchDescriptor bias_desc;
+      bias_desc.set_count(1)
+          .set_height(1)
+          .set_width(1)
+          .set_feature_map_count(params.output_descriptor.feature_map_count())
+          .set_layout(params.output_dl);
+      params.bias_descriptor.CloneFrom(bias_desc);
+      if (!params_ptr) {
+        print_memref(metadata, "metadata");
+        for (int i = 0; i < 5 * N - 4; ++i) {
+          TAO_VLOG(0) << "\t#" << i << ": " << metadata.data[i];
+        }
+        ctx->signalError(Context::FAILURE, "illegal conv op");
+        return;
+      }
+
+      if (!PickBestAlgorithm<T>(*params_ptr, operand_se_buffers, result_buffer,
+                                stream, ctx, gpu_driver)) {
+        ctx->signalError(Context::FAILURE, "fail to tune conv op");
+        return;
+      }
+
+      it = cache.emplace(key, std::move(*params_ptr)).first;
+    }
+    params = &it->second;
+  }
+
+  if (params == nullptr) {
+    ctx->signalError(Context::FAILURE,
+                     "unable to build conv params to launch a conv");
+    return;
+  }
+
+  ScratchAllocator scratch_allocator(ctx, gpu_driver);
+  Status launch_status = RunCudnnConvolution<T>(
+      *params, operand_se_buffers, result_buffer, &scratch_allocator, stream,
+      /*profile_result*/ nullptr);
+  if (!launch_status.ok()) {
+    std::string err_msg =
+        "d_conv launch failed: " + launch_status.error_message();
+    ctx->signalError(Context::FAILURE, err_msg);
+  }
+  return;
+}
+
 }  // namespace gpu_conv_impl
 
 ////////////////////////////////////////////////////////////////////////
@@ -1524,6 +1740,8 @@ TAO_RAL_API("ral_conv", "gpu",
 TAO_RAL_API("ral_conv", "gpu", gpu::se_impl::gpu_conv_impl::ral_conv<float, 3>);
 TAO_RAL_API("ral_conv", "gpu",
             gpu::se_impl::gpu_conv_impl::ral_conv<Eigen::half, 3>);
+TAO_RAL_API("ral_qconv", "gpu",
+            gpu::se_impl::gpu_conv_impl::ral_qconv<int8_t, 4>);
 
 }  // namespace ral
 }  // namespace tao


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_p_nhwc_i8_per_tensor.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_p_nhwc_i8_per_tensor.mlir
new file mode 100644
index 00000000000..6da05e796b3
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_p_nhwc_i8_per_tensor.mlir
@@ -0,0 +1,35 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%input: tensor<?x?x?x?xf32>) -> (tensor<?x?x?x16xf32>) 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<0.5> : 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?x?x?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?x?x?xf32>
+
+      %weight:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.0> : tensor<1x1x16x16xf32>} : () -> tensor<1x1x16x16xf32>
+      %weight_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.0> : tensor<16xf32>} : () -> tensor<16xf32>
+      %weight_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<16xi32>} : () -> tensor<16xi32>
+      %quantized_weight:2 = tf_executor.island wraps "tf.DiscFakeQuant"(%weight, %weight_scale, %weight_zero_point) {
+          axis = [3], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 127 : i64, quant_min = -128 : i64} :
+      (tensor<1x1x16x16xf32>, tensor<16xf32>, tensor<16xi32>) -> tensor<1x1x16x16xf32>
+
+      %result:2 = tf_executor.island wraps "tf.Conv2D"(%quantized_input, %quantized_weight)
+      {
+        data_format = "NHWC",
+        dilations = [1, 1, 1, 1],
+        padding = "SAME",
+        strides = [1, 1, 1, 1]
+      } : (tensor<?x?x?x?xf32>, tensor<1x1x16x16xf32>) -> tensor<?x?x?x16xf32>
+
+      %result_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.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<?x?x?x16xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?x?x16xf32>
+
+      tf_executor.fetch %quantized_result : tensor<?x?x?x16xf32>
+    }
+    return %graph : tensor<?x?x?x16xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_conv2d.cc b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_conv2d.cc
index 2fa31ecaba4..960689da510 100644
--- a/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_conv2d.cc
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_conv2d.cc
@@ -56,4 +56,25 @@ TEST(TFQuantziedConv2d, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_CHANNEL) {
       /*expect_output_vals*/ {output}));
 }
 
+// Comment qconv2d int8 GPU test since it only supports NVIDIA Ampere
+// architecture currently. Will add it back when ut running on machine whose GPU
+// architecture >= Ampere
+
+// TEST(TFQuantziedConv2d, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_TENSOR) {
+//   std::vector<float> inputs(1 * 56 * 56 * 16, 1);
+//   tensorflow::Tensor output(tensorflow::DataType::DT_FLOAT, {1, 56, 56, 16});
+//   auto datas = output.flat<float>();
+//   for (int i = 0; i < output.NumElements(); ++i) datas(i) = 16;
+//   EXPECT_TRUE(feature_test_main(
+//       /*mlir_file_path*/ c_ft_path +
+//           "quantized_conv2d_p_nhwc_i8_per_tensor.mlir",
+//       /*backend_types*/ {BackendType::kCuda},
+//       /*num_inputs*/ 1,
+//       /*num_outputs*/ 1,
+//       /*input_descriptors*/ {"1x56x56x16xf32_X"},
+//       /*output_descriptors*/ {"f32_X"},
+//       /*input_vals*/ {inputs},
+//       /*expect_output_vals*/ {output}));
+// }
+
 }  // namespace mlir_test