From f473dd295d241787b014423a9cbaf83fbe61a5fc Mon Sep 17 00:00:00 2001
From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com>
Date: Mon, 23 Nov 2020 12:04:32 -0800
Subject: [PATCH] Add QLinearMatMul(u8s8) (#5899)

---
 .../cpu/math/quantize_linear_matmul.cc        | 30 ++------------
 .../cpu/math/quantize_linear_matmul_test.cc   | 41 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc b/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc
index bd0add9ae9..9e1c1bc686 100644
--- a/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc
+++ b/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc
@@ -7,7 +7,6 @@
 #include "core/providers/common.h"
 #include "core/util/math_cpuonly.h"
 #include "core/util/qmath.h"
-#include "core/util/gemmlowp_common.h"
 #include "core/mlas/inc/mlas.h"
 
 namespace onnxruntime {
@@ -26,7 +25,7 @@ ONNX_OPERATOR_KERNEL_EX(
     kCpuExecutionProvider,
     KernelDefBuilder()
         .TypeConstraint("T1", DataTypeImpl::GetTensorType<uint8_t>())
-        .TypeConstraint("T2", DataTypeImpl::GetTensorType<uint8_t>())
+        .TypeConstraint("T2", {DataTypeImpl::GetTensorType<uint8_t>(), DataTypeImpl::GetTensorType<int8_t>()})
         .TypeConstraint("T3", DataTypeImpl::GetTensorType<uint8_t>()),
     QLinearMatMul);
 
@@ -70,32 +69,24 @@ Status QLinearMatMul::Compute(OpKernelContext* ctx) const {
 
   const float real_multiplier = (a_scale_data * b_scale_data) / y_scale_data;
 
-#ifdef MLAS_SUPPORTS_GEMM_U8X8
   AllocatorPtr alloc;
   ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc));
   auto gemm_output_data = alloc->Alloc(SafeInt<size_t>(sizeof(int32_t)) *
                                        static_cast<size_t>(helper.M()) * static_cast<size_t>(helper.N()));
   BufferUniquePtr gemm_output_buffer(gemm_output_data, BufferDeleter(alloc));
   auto* gemm_output = static_cast<int32_t*>(gemm_output_buffer.get());
-#else
-  // Compute the fixed point multiplier and shift for requantizing with GEMMLOWP.
-  int32_t integer_multiplier;
-  int right_shift;
-  QuantizeMultiplier(real_multiplier, &integer_multiplier, &right_shift);
-#endif
 
   for (size_t i = 0; i < helper.OutputOffsets().size(); i++) {
-#ifdef MLAS_SUPPORTS_GEMM_U8X8
     QGemm(static_cast<int>(helper.M()),
           static_cast<int>(helper.N()),
           static_cast<int>(helper.K()),
           a->template Data<uint8_t>() + helper.LeftOffsets()[i],
           static_cast<int>(helper.K()),
           *a_offset->template Data<uint8_t>(),
-          b->template Data<uint8_t>() + helper.RightOffsets()[i],
+          static_cast<const uint8_t*>(b->DataRaw()) + helper.RightOffsets()[i],
           static_cast<int>(helper.N()),
-          *b_offset->template Data<uint8_t>(),
-          false,
+          *static_cast<const uint8_t*>(b_offset->DataRaw()),
+          b->IsDataType<int8_t>(),
           gemm_output,
           static_cast<int>(helper.N()),
           ctx->GetOperatorThreadPool());
@@ -108,19 +99,6 @@ Status QLinearMatMul::Compute(OpKernelContext* ctx) const {
                          &real_multiplier,
                          false,
                          *y_offset->template Data<uint8_t>());
-#else
-    GemmlowpMultiplyu8u8_u8(a->template Data<uint8_t>() + helper.LeftOffsets()[i],
-                            b->template Data<uint8_t>() + helper.RightOffsets()[i],
-                            y->template MutableData<uint8_t>() + helper.OutputOffsets()[i],
-                            *a_offset->template Data<uint8_t>(),
-                            *b_offset->template Data<uint8_t>(),
-                            *y_offset->template Data<uint8_t>(),
-                            static_cast<int>(helper.M()),
-                            static_cast<int>(helper.N()),
-                            static_cast<int>(helper.K()),
-                            integer_multiplier,
-                            right_shift);
-#endif
   }
 
   return Status::OK();
diff --git a/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc b/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc
index d6b17cf5e9..0c95ff70e1 100644
--- a/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc
+++ b/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc
@@ -7,7 +7,7 @@
 namespace onnxruntime {
 namespace test {
 
-TEST(QuantizeLinearMatmulOpTest, QLinearMatMul3D) {
+TEST(QuantizeLinearMatmulOpTest, QLinearMatMul3D_U8U8) {
   OpTester test("QLinearMatMul", 10);
   test.AddInput<uint8_t>("T1", {2, 2, 4},
                          {208, 236, 0, 238,
@@ -45,6 +45,44 @@ TEST(QuantizeLinearMatmulOpTest, QLinearMatMul3D) {
   test.Run();
 }
 
+TEST(QuantizeLinearMatmulOpTest, QLinearMatMul3D_U8S8) {
+  OpTester test("QLinearMatMul", 10);
+  test.AddInput<uint8_t>("T1", {2, 2, 4},
+                         {208, 126, 0, 238,
+                          3, 214, 255, 29,
+
+                          208, 236, 0, 238,
+                          3, 214, 255, 29});
+
+  test.AddInput<float>("a_scale", {}, {0.0066f});
+  test.AddInput<uint8_t>("a_zero_point", {}, {113});
+
+  test.AddInput<int8_t>("T2", {2, 4, 3},
+                         {-43, 51, -34,
+                          60, 26, -17,
+                          0, 63, -55,
+                          47, -29, -31,
+
+                          -62, 51, -42,
+                          60, 26, -22,
+                          0, -8, -19,
+                          37, -2, -47});
+
+  test.AddInput<float>("b_scale", {}, {0.00802f});
+  test.AddInput<int8_t>("b_zero_point", {}, {-2});
+
+  test.AddInput<float>("y_scale", {}, {0.0123f});
+  test.AddInput<uint8_t>("y_zero_point", {}, {118});
+  test.AddOutput<uint8_t>("T3", {2, 2, 3},
+                          {130, 95, 114,
+                           148, 155, 105,
+
+                           146, 157, 75,
+                           160, 101, 134});
+
+  test.Run();
+}
+
 static void QLinearMatMul2DTest(bool only_t1_not_initializer) {
   // Test non-empty inputs
   OpTester test_non_empty("QLinearMatMul", 10);
@@ -83,5 +121,6 @@ TEST(QuantizeLinearMatmulOpTest, QLinearMatMul) {
 TEST(QuantizeLinearMatmulOpTest, QLinearMatMulAllInputExceptT1AreInitializers) {
   QLinearMatMul2DTest(true);
 }
+
 }  // namespace test
 }  // namespace onnxruntime