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 91a2c0f2638..be74c3194d8 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -222,6 +222,11 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(createCSEPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
 
+  // quantization related passes.
+  pm.addNestedPass<FuncOp>(disc_ral::createDiscConvertFakeQuantOpPass());
+  pm.addNestedPass<FuncOp>(
+      disc_ral::createDiscLowerQuantizeAndDequantizePass());
+
   bool enable_shape_constraint_ir = useShapeConstraintIR();
   if (!enable_shape_constraint_ir) {
     // propagate some known shape information.
diff --git a/tao_compiler/mlir/xla/ral/context/common_context_impl_quantization.cc b/tao_compiler/mlir/xla/ral/context/common_context_impl_quantization.cc
index 6802d00d92a..f7b2579ad45 100644
--- a/tao_compiler/mlir/xla/ral/context/common_context_impl_quantization.cc
+++ b/tao_compiler/mlir/xla/ral/context/common_context_impl_quantization.cc
@@ -159,10 +159,184 @@ void ral_qconv_s8_s8_s8(
   }
 }
 
+template <int NDims>
+void ral_qconv_acl_s8_s8_s8_per_channel(
+    ExecutionContext* ctx, opaque_t /*stream_handle*/,
+    MemRefType<int8_t, NDims> input, MemRefType<int8_t, NDims> weight,
+    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<int8_t, NDims> result,
+    MemRefType<int32_t, 1> metadata) {
+  CpuTimer timer("ral_qconv_acl_s8_s8_s8_per_channel");
+  if (isEmptyMemref(input) || isEmptyMemref(weight) || isEmptyMemref(result)) {
+    TAO_VLOG(1)
+        << "ral_qconv_acl_s8_s8_s8_per_channel: early return for empty tensor";
+    return;
+  }
+  ConvParams params;
+  if (!parseConvParams(ctx, input, weight, padding, result, metadata,
+                       &params)) {
+    ctx->signalError(Context::FAILURE, "invalid conv params");
+  }
+
+  if (TAO_VLOG_IS_ON(2)) {
+    for (int i = 0; i < Size(input); ++i) {
+      TAO_VLOG(0) << "input[" << i
+                  << "] = " << static_cast<int32_t>(input.data[i]);
+    }
+    for (int i = 0; i < Size(weight); ++i) {
+      TAO_VLOG(0) << "weight[" << i
+                  << "] = " << static_cast<int32_t>(weight.data[i]);
+    }
+  }
+
+  if (TAO_VLOG_IS_ON(1)) {
+    TAO_VLOG(0) << "input scale = " << inputScales.data[0];
+    TAO_VLOG(0) << "input zero point = " << inputZeroPoints.data[0];
+    TAO_VLOG(0) << "result scale = " << resultScales.data[0];
+    TAO_VLOG(0) << "result zero point = " << resultZeroPoints.data[0];
+    for (int i = 0; i < weightScales.sizes[0]; ++i)
+      TAO_VLOG(0) << "weight_scale[" << i << "] = " << weightScales.data[i];
+    for (int i = 0; i < weightZeroPoints.sizes[0]; ++i)
+      TAO_VLOG(0) << "weight_zero_point[" << i
+                  << "] = " << weightZeroPoints.data[i];
+  }
+
+  if (params.groups > 1) {
+    ctx->signalError(Context::FAILURE, "invalid conv params");
+  }
+
+  auto src_dims = params.src.get_dims();
+  auto dst_dims = params.dst.get_dims();
+  auto weight_dims = params.weight.get_dims();
+  int N = src_dims[0];
+  int Ci = src_dims[1];
+  int Ih = src_dims[2];
+  int Iw = src_dims[3];
+  int Co = dst_dims[1];
+  int Oh = dst_dims[2];
+  int Ow = dst_dims[3];
+  int Kh = weight_dims[2];
+  int Kw = weight_dims[3];
+
+  if (TAO_VLOG_IS_ON(1)) {
+    TAO_VLOG(1) << "N = " << N;
+    TAO_VLOG(1) << "Ih = " << Ih;
+    TAO_VLOG(1) << "Iw = " << Iw;
+    TAO_VLOG(1) << "Ci = " << Ci;
+    TAO_VLOG(1) << "Oh = " << Oh;
+    TAO_VLOG(1) << "Ow = " << Ow;
+    TAO_VLOG(1) << "Co = " << Co;
+    TAO_VLOG(1) << "Kh = " << Kh;
+    TAO_VLOG(1) << "Kw = " << Kw;
+    TAO_VLOG(0) << "params.strides[1] = " << params.strides[1];
+    TAO_VLOG(0) << "params.strides[0] = " << params.strides[0];
+    TAO_VLOG(0) << "params.padding_l[1] = " << params.padding_l[1];
+    TAO_VLOG(0) << "params.padding_l[0] = " << params.padding_l[0];
+    TAO_VLOG(0) << "params.padding_r[1] = " << params.padding_r[1];
+    TAO_VLOG(0) << "params.padding_r[0] = " << params.padding_r[0];
+    TAO_VLOG(0) << "params.dilates[1] = " << params.dilates[1];
+    TAO_VLOG(0) << "params.dilates[0] = " << params.dilates[0];
+  }
+
+  auto AclQconvCreator = [&](const arm_compute::ITensorPack* pack) {
+    std::shared_ptr<AclConvInfo> info(new AclConvInfo);
+    auto src_shape = TensorShape(Ci, Iw, Ih, N);
+    auto weights_shape = TensorShape(Ci, Kw, Kh, Co);
+    auto dst_shape = TensorShape(Co, Ow, Oh, N);
+
+    DataLayout data_layout = DataLayout::NHWC;
+    DataType data_type = DataType::QASYMM8_SIGNED;
+    TensorInfo src_info = TensorInfo(src_shape, 1, data_type, data_layout);
+    TensorInfo weights_info =
+        TensorInfo(weights_shape, 1, DataType::QSYMM8_PER_CHANNEL, data_layout);
+    TensorInfo dst_info = TensorInfo(dst_shape, 1, data_type, data_layout);
+
+    const QuantizationInfo src_qinfo = QuantizationInfo();
+    src_info.set_quantization_info(
+        QuantizationInfo(*inputScales.data, *inputZeroPoints.data));
+    std::vector<float> scales(weightScales.data,
+                              weightScales.data + weightScales.sizes[0]);
+    std::vector<int32_t> zero_points(
+        weightZeroPoints.data,
+        weightZeroPoints.data + weightZeroPoints.sizes[0]);
+    weights_info.set_quantization_info(
+        QuantizationInfo(std::move(scales), std::move(zero_points)));
+    dst_info.set_quantization_info(
+        QuantizationInfo(*resultScales.data, *resultZeroPoints.data));
+
+    info->src.allocator()->init(src_info);
+    info->weights.allocator()->init(weights_info);
+    info->dst.allocator()->init(dst_info);
+    info->src.allocator()->import_memory(input.data);
+    info->weights.allocator()->import_memory(weight.data);
+    info->dst.allocator()->import_memory(result.data);
+
+    if (!info->op.validate(
+            &src_info, &weights_info, nullptr, &dst_info,
+            PadStrideInfo{params.strides[1], params.strides[0],
+                          params.padding_l[1], params.padding_r[1],
+                          params.padding_l[0], params.padding_r[0],
+                          DimensionRoundingType::FLOOR},
+            WeightsInfo{}, Size2D{params.dilates[1], params.dilates[0]})) {
+      ctx->signalError(
+          Context::FAILURE,
+          "fail to validate ral_qconv_acl_s8_s8_s8_per_channel conv");
+    } else {
+      info->op.configure(&info->src, &info->weights, nullptr, &info->dst,
+                         PadStrideInfo{params.strides[1], params.strides[0],
+                                       params.padding_l[1], params.padding_r[1],
+                                       params.padding_l[0], params.padding_r[0],
+                                       DimensionRoundingType::FLOOR},
+                         WeightsInfo{},
+                         Size2D{params.dilates[1], params.dilates[0]});
+    }
+    if (pack) info->op.reuse_packed_weight(*pack);
+    info->op.prepare(&info->src, &info->weights, nullptr, &info->dst);
+    return info;
+  };
+
+  std::shared_ptr<AclConvInfo> info;
+  if (isWeightPrePackingEnabled() && params.weight_is_const) {
+    std::string unique_name = "disc.ral_qconv_acl_s8_s8_s8_per_channel";
+    auto state = ctx->getOrCreateResource<AclConvState>(
+        unique_name, []() { return new AclConvState; });
+    auto key = makeConvParamsKey(input, weight, padding, result, metadata,
+                                 kDiscCpuDefaultThreadId);
+    info = state->getOrCreate(key, AclQconvCreator);
+  } else {
+    info = AclQconvCreator(nullptr);
+  }
+
+  // TOOD(disc): re-import quantization info when online-quantization is
+  // supported.
+  info->src.allocator()->import_memory(input.data);
+  info->dst.allocator()->import_memory(result.data);
+  info->op.run(&info->src, &info->weights, nullptr, &info->dst);
+
+  if (TAO_VLOG_IS_ON(2)) {
+    for (int i = 0; i < Size(result); ++i) {
+      TAO_VLOG(0) << "result[" << i
+                  << "] = " << static_cast<int32_t>(result.data[i]);
+    }
+  }
+
+  timer.Stop();
+  if (isProfilingEnabled()) {
+    dumpConvLikeKernelProflingInfo<int8_t>(
+        params, timer.GetNanoSeconds(), "ral_qconv_acl_s8_s8_s8_per_channel");
+  }
+}
+
 }  // namespace
 
+// Deprecated, remove such implementation once refactor is done.
 TAO_RAL_API("ral_qconv_s8_s8_s8", "cpu", ral_qconv_s8_s8_s8<4>);
 
+// new fake_quant based implementation.
+TAO_RAL_API("ral_qconv", "cpu", ral_qconv_acl_s8_s8_s8_per_channel<4>);
+
 }  // namespace ral
 }  // namespace tao
 #endif


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_p_nhwc_i8_per_channel.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_p_nhwc_i8_per_channel.mlir
new file mode 100644
index 00000000000..2792a4486c2
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_p_nhwc_i8_per_channel.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?x8xf32>) 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<8.333333e-3> : 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<-0.8> : tensor<1x1x25x8xf32>} : () -> tensor<1x1x25x8xf32>
+      %weight_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.568627e-2> : tensor<8xf32>} : () -> tensor<8xf32>
+      %weight_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<8xi32>} : () -> tensor<8xi32>
+      %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<1x1x25x8xf32>, tensor<8xf32>, tensor<8xi32>) -> tensor<1x1x25x8xf32>
+
+      %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<1x1x25x8xf32>) -> tensor<?x?x?x8xf32>
+
+      %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?x8xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?x?x8xf32>
+
+      tf_executor.fetch %quantized_result : tensor<?x?x?x8xf32>
+    }
+    return %graph : tensor<?x?x?x8xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_s_nhwc_i8_per_channel.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_s_nhwc_i8_per_channel.mlir
new file mode 100644
index 00000000000..b9cb156e7a3
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_conv2d_s_nhwc_i8_per_channel.mlir
@@ -0,0 +1,35 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%input: tensor<1x2x2x25xf32>) -> (tensor<1x2x2x8xf32>) 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<8.333333e-3> : 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<1x2x2x25xf32>, tensor<f32>, tensor<i32>) -> tensor<1x2x2x25xf32>
+
+      %weight:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.8> : tensor<1x1x25x8xf32>} : () -> tensor<1x1x25x8xf32>
+      %weight_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.568627e-2> : tensor<8xf32>} : () -> tensor<8xf32>
+      %weight_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<8xi32>} : () -> tensor<8xi32>
+      %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<1x1x25x8xf32>, tensor<8xf32>, tensor<8xi32>) -> tensor<1x1x25x8xf32>
+
+      %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<1x2x2x25xf32>, tensor<1x1x25x8xf32>) -> tensor<1x2x2x8xf32>
+
+      %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<1x2x2x8xf32>, tensor<f32>, tensor<i32>) -> tensor<1x2x2x8xf32>
+
+      tf_executor.fetch %quantized_result : tensor<1x2x2x8xf32>
+    }
+    return %graph : tensor<1x2x2x8xf32>
+  }
+}
\ 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
new file mode 100644
index 00000000000..e1fb587b7b8
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_conv2d.cc
@@ -0,0 +1,59 @@
+/* Copyright 2022 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/tests/mlir_feature_test.h"
+#include "tensorflow/compiler/mlir/disc/tests/mlir_test.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir_test {
+
+const std::string c_ft_path =
+    "tensorflow/compiler/mlir/disc/tests/tensorflow_ops/data/";
+
+TEST(TFQuantziedConv2d, STATIC_SHAPE_NHWC_I8_PER_TENSOR) {
+  std::vector<float> inputs(1 * 2 * 2 * 25, -0.6);
+  tensorflow::Tensor output(tensorflow::DataType::DT_FLOAT, {1, 2, 2, 8});
+  auto datas = output.flat<float>();
+  for (int i = 0; i < output.NumElements(); ++i) datas(i) = 12.0;
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path +
+          "quantized_conv2d_s_nhwc_i8_per_channel.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"1x2x2x25xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {inputs},
+      /*expect_output_vals*/ {output}));
+}
+
+TEST(TFQuantziedConv2d, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_TENSOR) {
+  std::vector<float> inputs(1 * 2 * 2 * 25, -0.6);
+  tensorflow::Tensor output(tensorflow::DataType::DT_FLOAT, {1, 2, 2, 8});
+  auto datas = output.flat<float>();
+  for (int i = 0; i < output.NumElements(); ++i) datas(i) = 12.0;
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path +
+          "quantized_conv2d_p_nhwc_i8_per_channel.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"1x2x2x25xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {inputs},
+      /*expect_output_vals*/ {output}));
+}
+
+}  // namespace mlir_test