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/transforms/disc_lower_to_library_call.cc b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
index 254f26f52b7..fe62438dd1f 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
@@ -65,6 +65,7 @@ using lmhlo::ReshapeOp;
 using lmhlo_disc::CustomCallOp;
 using lmhlo_disc::D2HOp;
 using lmhlo_disc::H2DOp;
+using lmhlo_disc::QuantizedDotGeneralOp;
 using lmhlo_disc::QuantizedDynamicConvOp;
 
 // Suppose that the first argument of the function is the ctx value
@@ -179,10 +180,14 @@ struct GpuCopyOpConvertor : public OpRewritePattern<OpTy> {
   StringRef target_;
 };
 
-struct DotGeneralOpConvertor : public OpRewritePattern<DotGeneralOp> {
-  using OpRewritePattern<DotGeneralOp>::OpRewritePattern;
+template <typename OpTy>
+struct DotGeneralLikeConverter : public OpRewritePattern<OpTy> {
+  DotGeneralLikeConverter(MLIRContext* context, StringRef target)
+      : OpRewritePattern<OpTy>::OpRewritePattern(context) {
+    this->target_ = target;
+  }
 
-  LogicalResult matchAndRewrite(DotGeneralOp op,
+  LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter& rewriter) const override {
     Value ctx = GetContextValueFromFunctionArguments(op);
     if (!ctx) {
@@ -204,7 +209,7 @@ struct DotGeneralOpConvertor : public OpRewritePattern<DotGeneralOp> {
       return op.emitOpError() << "unmatched batch dims size.";
     }
 
-    int rank = op.getOperand(0).getType().cast<MemRefType>().getRank();
+    int rank = op.getOperand(0).getType().template cast<MemRefType>().getRank();
     for (auto&& z : llvm::zip(lhs_batching_dims, rhs_batching_dims)) {
       if ((std::get<0>(z) >= rank - 2) || (std::get<1>(z) >= rank - 2)) {
         return op.emitOpError() << "unsupported batch dims.";
@@ -242,11 +247,14 @@ struct DotGeneralOpConvertor : public OpRewritePattern<DotGeneralOp> {
 
     bool on_gpu = placement_utils::isGpuMemRef(op->getOperand(2));
     rewriter.replaceOpWithNewOp<DispatchOp>(op, llvm::None, ctx, newOperands,
-                                            "ral_gemm", false,
+                                            target_, false,
                                             on_gpu ? "gpu" : "cpu");
 
     return success();
   }
+
+ private:
+  StringRef target_;
 };
 
 template <typename OpTy>
@@ -619,7 +627,6 @@ struct DiscLowerToLibraryCallPass
       CopyLikeOpConvertor<DynamicReshapeOp>,
       CopyLikeOpConvertor<ReshapeOp>,
       CustomCallOpConvertor,
-      DotGeneralOpConvertor,
       RecvInputOpConvertor,
       ConvConverter,
       SendOutputOpConvertor
@@ -633,6 +640,9 @@ struct DiscLowerToLibraryCallPass
                                                              "ral_conv");
     patterns.insert<DynamicConvLikeConverter<QuantizedDynamicConvOp>>(
         context, "ral_qconv");
+    patterns.insert<DotGeneralLikeConverter<DotGeneralOp>>(context, "ral_gemm");
+    patterns.insert<DotGeneralLikeConverter<QuantizedDotGeneralOp>>(
+        context, "ral_qgemm");
 
     if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
       func.emitError("applyPatternsAndFoldGreedily does not converge");


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-lower-to-library-call.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-lower-to-library-call.mlir
index 219bd0071fd..a14166a0b80 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-lower-to-library-call.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-lower-to-library-call.mlir
@@ -1,4 +1,4 @@
-// RUN: disc-opt -disc-lower-to-library-call %s -o - | FileCheck %s
+// RUN: disc-opt -disc-lower-to-library-call --split-input-file %s -o - | FileCheck %s
 
 // CHECK-LABEL: func.func @test_recv_input_and_send_output
 // CHECK-SAME: (%[[CTX:.*]]: !disc_ral.context) {
@@ -486,3 +486,43 @@ func.func @quantized_dynamic_conv(%arg0: !disc_ral.context) {
   } : (memref<?x?x?x?xi8, "gpu">, memref<?x?x?x?xi8, "gpu">, memref<4xi32, "cpu">, memref<f32, "gpu">, memref<i32, "gpu">, memref<f32, "gpu">, memref<i32, "gpu">, memref<f32, "gpu">, memref<i32, "gpu">, memref<?x?x?x?xi8, "gpu">) -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: quantized_dot_general
+func.func @quantized_dot_general(%arg0: !disc_ral.context) {
+  %c0 = arith.constant 0 : index
+  %0 = "disc_ral.recv_input"(%arg0, %c0) : (!disc_ral.context, index) -> memref<?x?x?x?xi8, "gpu">
+  %c1 = arith.constant 1 : index
+  %1 = "disc_ral.recv_input"(%arg0, %c1) : (!disc_ral.context, index) -> memref<?x?x?x?xi8, "gpu">
+  %c2 = arith.constant 2 : index
+  %2 = "disc_ral.recv_input"(%arg0, %c2) : (!disc_ral.context, index) -> memref<f32, "gpu">
+  %c3 = arith.constant 3 : index
+  %3 = "disc_ral.recv_input"(%arg0, %c3) : (!disc_ral.context, index) -> memref<i32, "gpu">
+  %c4 = arith.constant 4 : index
+  %4 = "disc_ral.recv_input"(%arg0, %c4) : (!disc_ral.context, index) -> memref<f32, "gpu">
+  %c5 = arith.constant 5 : index
+  %5 = "disc_ral.recv_input"(%arg0, %c5) : (!disc_ral.context, index) -> memref<i32, "gpu">
+  %c6 = arith.constant 6 : index
+  %6 = "disc_ral.recv_input"(%arg0, %c6) : (!disc_ral.context, index) -> memref<f32, "gpu">
+  %c7 = arith.constant 7 : index
+  %7 = "disc_ral.recv_input"(%arg0, %c7) : (!disc_ral.context, index) -> memref<i32, "gpu">
+  %c8 = arith.constant 8 : index
+  %8 = "disc_ral.recv_input"(%arg0, %c8) : (!disc_ral.context, index) -> memref<?x?x?x?xi8, "gpu">
+  // CHECK: %[[DEFAULT_STREAM:.*]] = llvm.inttoptr {{.*}} : i32 to !llvm.ptr<i8>
+  // CHECK: disc_ral.dispatch
+  // CHECK-SAME: backend_config = "gpu"
+  // CHECK-SAME: call_target_name = "ral_qgemm"
+  "lmhlo_disc.quantized_dot_general"(%0, %1, %2, %3, %4, %5, %6, %7, %8) {
+    axis = dense<> : tensor<0xi64>,
+    use_dynamic = false,
+    use_symmetric = true,
+    dot_dimension_numbers = #mhlo.dot<
+      lhs_batching_dimensions = [0, 1],
+      rhs_batching_dimensions = [0, 1],
+      lhs_contracting_dimensions = [3],
+      rhs_contracting_dimensions = [3]
+    >
+  } : (memref<?x?x?x?xi8, "gpu">, memref<?x?x?x?xi8, "gpu">, memref<f32, "gpu">, memref<i32, "gpu">, memref<f32, "gpu">, memref<i32, "gpu">, memref<f32, "gpu">, memref<i32, "gpu">, memref<?x?x?x?xi8, "gpu">) -> ()
+  return
+}