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/xla/ral/context/common_context_impl_acl.cc b/tao_compiler/mlir/xla/ral/context/common_context_impl_acl.cc
index b10ea5b87b6..0a526d6e865 100644
--- a/tao_compiler/mlir/xla/ral/context/common_context_impl_acl.cc
+++ b/tao_compiler/mlir/xla/ral/context/common_context_impl_acl.cc
@@ -25,6 +25,7 @@
 #include "src/common/utils/Log.h"
 #include "src/cpu/operators/CpuConv2d.h"
 #include "src/cpu/operators/CpuDepthwiseConv2d.h"
+#include "src/cpu/operators/CpuGemmLowpMatrixMultiplyCore.h"
 
 namespace arm_compute {
 using namespace arm_compute::experimental;
@@ -628,6 +629,122 @@ void DISCNEDepthwiseConvolutionLayer::prepare(ITensor* input,
   }
 }
 
+//////////////////// NEGEMMLowpMatrixMultiplyCore /////////////////
+
+struct DISCNEGEMMLowpMatrixMultiplyCore::Impl {
+  const ITensor* b{nullptr};
+  std::unique_ptr<cpu::CpuGemmLowpMatrixMultiplyCore> op{nullptr};
+  // persistent tensors are usually those buffers that are used to store packed
+  // weights.
+  ITensorPack persistent_pack{};
+  MemoryRequirements aux_mem_req{};
+  WorkspaceData<Tensor> workspace{};
+  bool is_prepared{false};
+};
+
+DISCNEGEMMLowpMatrixMultiplyCore::DISCNEGEMMLowpMatrixMultiplyCore(
+    std::shared_ptr<IMemoryManager> memory_manager,
+    IWeightsManager* weights_manager)
+    : _impl(std::make_unique<Impl>()) {}
+DISCNEGEMMLowpMatrixMultiplyCore::~DISCNEGEMMLowpMatrixMultiplyCore() = default;
+
+void DISCNEGEMMLowpMatrixMultiplyCore::configure(const ITensor* a,
+                                                 const ITensor* b,
+                                                 const ITensor* c,
+                                                 ITensor* output,
+                                                 const GEMMInfo& gemm_info) {
+  ARM_COMPUTE_ERROR_ON_NULLPTR(a, b, output);
+  _impl->b = b;
+  _impl->op = std::make_unique<cpu::CpuGemmLowpMatrixMultiplyCore>();
+  _impl->op->configure(a->info(), b->info(),
+                       (c != nullptr ? c->info() : nullptr), output->info(),
+                       gemm_info);
+  _impl->aux_mem_req = _impl->op->workspace();
+}
+
+Status DISCNEGEMMLowpMatrixMultiplyCore::validate(const ITensorInfo* a,
+                                                  const ITensorInfo* b,
+                                                  const ITensorInfo* c,
+                                                  const ITensorInfo* output,
+                                                  const GEMMInfo& gemm_info) {
+  return cpu::CpuGemmLowpMatrixMultiplyCore::validate(a, b, c, output,
+                                                      gemm_info);
+}
+
+void DISCNEGEMMLowpMatrixMultiplyCore::run() {
+  ARM_COMPUTE_ERROR("Not supported.");
+}
+
+void DISCNEGEMMLowpMatrixMultiplyCore::prepare() {
+  ARM_COMPUTE_ERROR("Not supported.");
+}
+
+void DISCNEGEMMLowpMatrixMultiplyCore::run(ITensor* a, const ITensor* b,
+                                           const ITensor* c, ITensor* output) {
+  prepare(a, b, c, output);
+
+  ITensorPack run_pack = {{TensorType::ACL_SRC_0, a},
+                          {TensorType::ACL_SRC_1, b},
+                          {TensorType::ACL_SRC_2, c},
+                          {TensorType::ACL_DST, output}};
+  DISCAllocator allocator;
+  DISCBuffers buffers;
+  auto workspace = manage_runtime_workspace<Tensor>(
+      _impl->aux_mem_req, run_pack, _impl->persistent_pack, allocator, buffers);
+  _impl->op->run(run_pack);
+}
+
+void DISCNEGEMMLowpMatrixMultiplyCore::prepare(ITensor* a, const ITensor* b,
+                                               const ITensor* c,
+                                               ITensor* output) {
+  if (!_impl->is_prepared) {
+    ITensorPack prep_pack = {{TensorType::ACL_SRC_0, a},
+                             {TensorType::ACL_SRC_1, b},
+                             {TensorType::ACL_SRC_2, c},
+                             {TensorType::ACL_DST, output}};
+
+    std::vector<Tensor*> temporary_tensors;
+    for (const auto& req : _impl->aux_mem_req) {
+      if (req.size == 0) continue;
+      // re-use from exsiting buffer whenever possible.
+      if (auto aux_tensor = _impl->persistent_pack.get_tensor(req.slot)) {
+        prep_pack.add_tensor(req.slot, aux_tensor);
+        continue;
+      }
+
+      const auto aux_info = TensorInfo{TensorShape(req.size), 1, DataType::U8};
+      _impl->workspace.emplace_back(WorkspaceDataElement<Tensor>{
+          req.slot, req.lifetime, std::make_unique<Tensor>()});
+      auto aux_tensor = _impl->workspace.back().tensor.get();
+      ARM_COMPUTE_ERROR_ON_NULLPTR(aux_tensor);
+      aux_tensor->allocator()->init(aux_info, req.alignment);
+      aux_tensor->allocator()->allocate();
+      prep_pack.add_tensor(req.slot, aux_tensor);
+      if (req.lifetime == experimental::MemoryLifetime::Persistent) {
+        _impl->persistent_pack.add_tensor(req.slot, aux_tensor);
+      } else {
+        temporary_tensors.push_back(aux_tensor);
+      }
+    }
+
+    _impl->op->prepare(prep_pack);
+
+    for (auto aux_tensor : temporary_tensors) {
+      aux_tensor->allocator()->free();
+    }
+    _impl->is_prepared = true;
+  }
+}
+
+const ITensorPack& DISCNEGEMMLowpMatrixMultiplyCore::get_packed_weight() {
+  return _impl->persistent_pack;
+}
+
+void DISCNEGEMMLowpMatrixMultiplyCore::reuse_packed_weight(
+    const ITensorPack& pack) {
+  _impl->persistent_pack = pack;
+}
+
 }  // namespace arm_compute
 
 #endif  // defined(TAO_CPU_ONLY) && defined(TAO_AARCH64)
diff --git a/tao_compiler/mlir/xla/ral/context/common_context_impl_acl.h b/tao_compiler/mlir/xla/ral/context/common_context_impl_acl.h
index 85b292d5e0b..46fea94db68 100644
--- a/tao_compiler/mlir/xla/ral/context/common_context_impl_acl.h
+++ b/tao_compiler/mlir/xla/ral/context/common_context_impl_acl.h
@@ -578,6 +578,101 @@ class DISCNEDepthwiseConvolutionLayer : public IFunction {
   std::unique_ptr<Impl> _impl;
 };
 
+class DISCNEGEMMLowpMatrixMultiplyCore : public IFunction {
+ public:
+  /** Constructor */
+  DISCNEGEMMLowpMatrixMultiplyCore(
+      std::shared_ptr<IMemoryManager> memory_manager = nullptr,
+      IWeightsManager* weights_manager = nullptr);
+  /** Prevent instances of this class from being copied (As this class contains
+   * pointers) */
+  DISCNEGEMMLowpMatrixMultiplyCore(const DISCNEGEMMLowpMatrixMultiplyCore&) =
+      delete;
+  /** Default move constructor */
+  DISCNEGEMMLowpMatrixMultiplyCore(DISCNEGEMMLowpMatrixMultiplyCore&&) =
+      default;
+  /** Prevent instances of this class from being copied (As this class contains
+   * pointers) */
+  DISCNEGEMMLowpMatrixMultiplyCore& operator=(
+      const DISCNEGEMMLowpMatrixMultiplyCore&) = delete;
+  /** Default move assignment operator */
+  DISCNEGEMMLowpMatrixMultiplyCore& operator=(
+      DISCNEGEMMLowpMatrixMultiplyCore&&) = default;
+  /** Default destructor */
+  ~DISCNEGEMMLowpMatrixMultiplyCore();
+  /** Initialise the kernel's inputs, output
+   *
+   * Valid data layouts:
+   * - NHWC
+   * - NCHW
+   *
+   * Valid data type configurations:
+   * |src0           |src1               |src2     |dst            |
+   * |:--------------|:------------------|:--------|:--------------|
+   * |QASYMM8        |QASYMM8            |S32      |QASYMM8        |
+   * |QASYMM8        |QSYMM8_PER_CHANNEL |S32      |QASYMM8        |
+   * |QASYMM8        |QSYMM8             |S32      |QASYMM8        |
+   * |QASYMM8        |QASYMM8            |S32      |S32            |
+   * |QASYMM8        |QSYMM8_PER_CHANNEL |S32      |S32            |
+   * |QASYMM8        |QSYMM8             |S32      |S32            |
+   * |QASYMM8_SIGNED |QASYMM8_SIGNED     |S32      |QASYMM8_SIGNED |
+   * |QASYMM8_SIGNED |QSYMM8_PER_CHANNEL |S32      |QASYMM8_SIGNED |
+   * |QASYMM8_SIGNED |QSYMM8             |S32      |QASYMM8_SIGNED |
+   * |QASYMM8_SIGNED |QASYMM8_SIGNED     |S32      |S32            |
+   * |QASYMM8_SIGNED |QSYMM8_PER_CHANNEL |S32      |S32            |
+   * |QASYMM8_SIGNED |QSYMM8             |S32      |S32            |
+   *
+   * @note GEMM_LOWP:  low precision GEMM kernel
+   *  This kernel performs the following computations:
+   *
+   *  -# Convert a values from QASYMM8 to int32 and add a_offset to each of
+   * them.
+   *  -# Convert b values from QASYMM8 to int32 add b_offset to each of them.
+   *  -# Compute the matrix product of the resulting a * b in int32.
+   *
+   * @note The @p output type is S32 if @p gemm_info.type ==
+   * GEMMLowpOutputStageType::NONE. It is QASYMM8/QASYMM8_SIGNED otherwise
+   *
+   * @param[in]  a         First input tensor  (Matrix A). Data type supported:
+   * QASYMM8/QASYMM8_SIGNED.
+   * @param[in]  b         Second input tensor (Matrix B). Data type supported:
+   * QASYMM8/QASYMM8_SIGNED/QSYMM8/QSYMM8_PER_CHANNEL.
+   * @param[in]  c         Third input tensor  (Matrix C). It can be a nullptr.
+   * Data type supported: S32
+   * @param[out] output    Output tensor. Data type supported: Data type
+   * supported: S32/QASYMM8/QASYMM8_SIGNED
+   * @param[in]  gemm_info (Optional) Specifies if the matrix A and/or matrix B
+   * have been reshaped and if the reshape of matrix B should be executed only
+   * for the first run
+   */
+  void configure(const ITensor* a, const ITensor* b, const ITensor* c,
+                 ITensor* output, const GEMMInfo& gemm_info = GEMMInfo());
+  /** Static function to check if given info will lead to a valid configuration
+   * of @ref NEGEMMLowpMatrixMultiplyCore
+   *
+   * Similar to @ref NEGEMMLowpMatrixMultiplyCore::configure()
+   *
+   * @return a status
+   */
+  static Status validate(const ITensorInfo* a, const ITensorInfo* b,
+                         const ITensorInfo* c, const ITensorInfo* output,
+                         const GEMMInfo& gemm_info = GEMMInfo());
+
+  // Inherited methods overridden
+  void run() override;
+  void prepare() override;
+
+  void run(ITensor* a, const ITensor* b, const ITensor* c, ITensor* output);
+  void prepare(ITensor* a, const ITensor* b, const ITensor* c, ITensor* output);
+
+  const ITensorPack& get_packed_weight();
+  void reuse_packed_weight(const ITensorPack& pack);
+
+ private:
+  struct Impl;
+  std::unique_ptr<Impl> _impl;
+};
+
 }  // namespace arm_compute
 
 #endif  // defined(TAO_CPU_ONLY) && defined(TAO_AARCH64)
diff --git a/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.cc b/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.cc
index 0e37376591d..41bb3c2cb3d 100644
--- a/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.cc
+++ b/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.cc
@@ -92,6 +92,8 @@ int initWeightPrePackingCacheCapacity() {
   return std::atoi(env);
 }
 
+}  // namespace
+
 #if defined(TAO_AARCH64)
 bool applyACLThreadPoolConfigIfNotSet() {
   static bool status = []() {
@@ -118,8 +120,6 @@ bool applyACLThreadPoolConfigIfNotSet() {
 }
 #endif
 
-}  // namespace
-
 const std::thread::id kDiscCpuDefaultThreadId{};
 
 DiscCpuMathKernelMode GetDiscCpuMathKernelMode() {
@@ -559,46 +559,13 @@ TAO_RAL_API("ral_conv", "cpu", ral_conv<float, 3>);
 TAO_RAL_API("ral_conv", "cpu", ral_conv<float, 4>);
 TAO_RAL_API("ral_conv", "cpu", ral_conv<float, 5>);
 
-struct CpuGemmKey {
-  int m = -1;
-  int n = -1;
-  int k = -1;
-  int batch = 1;
-  bool transpose_a = false;
-  bool transpose_b = false;
-  opaque_t const_weight_ptr = nullptr;
-  // We need this since the matmul primitive is not thead safe.
-  // To enable large parallelism, we cache the primitive per-thread.
-  std::thread::id tid;
-
-  bool operator==(const CpuGemmKey& rhs) const {
-    return m == rhs.m && n == rhs.n && k == rhs.k && batch == rhs.batch &&
-           transpose_a == rhs.transpose_a && transpose_b == rhs.transpose_b &&
-           const_weight_ptr == rhs.const_weight_ptr && tid == rhs.tid;
-  }
-};
-
-struct CpuGemmKeyHasher {
-  std::size_t operator()(const CpuGemmKey& key) const {
-    std::size_t seed = std::hash<int>()(key.m);
-    hash_combine(seed, key.n);
-    hash_combine(seed, key.k);
-    hash_combine(seed, key.batch);
-    hash_combine(seed, key.transpose_a);
-    hash_combine(seed, key.transpose_b);
-    hash_combine(seed, key.const_weight_ptr);
-    hash_combine(seed, key.tid);
-    return seed;
-  }
-};
-
 template <typename TKey, typename TValue>
-using CpuGemmKeyMap = std::unordered_map<TKey, TValue, CpuGemmKeyHasher>;
+using CpuGemmKeyMap = std::unordered_map<TKey, TValue, typename TKey::Hasher>;
 
 #if defined(TAO_X86)
 class MklPackedWeight {
  public:
-  MklPackedWeight(ExecutionContext* ctx, const CpuGemmKey& key);
+  MklPackedWeight(ExecutionContext* ctx, const GEMMParamsKey& key);
   ~MklPackedWeight();
 
   opaque_t packed_weight() const { return packed_weight_; }
@@ -609,7 +576,8 @@ class MklPackedWeight {
   opaque_t packed_weight_ = nullptr;
 };
 
-MklPackedWeight::MklPackedWeight(ExecutionContext* ctx, const CpuGemmKey& key)
+MklPackedWeight::MklPackedWeight(ExecutionContext* ctx,
+                                 const GEMMParamsKey& key)
     : ctx_(ctx->getContext()),
       driver_(ctx->getDriver<cpu::CPUDriver>(cpu::CPUDriver::name())) {
   size_t packed_weight_size =
@@ -627,7 +595,7 @@ MklPackedWeight::~MklPackedWeight() {
 }
 
 using MklGemmCache =
-    ideep::utils::lru_cache<CpuGemmKey, std::shared_ptr<MklPackedWeight>,
+    ideep::utils::lru_cache<GEMMParamsKey, std::shared_ptr<MklPackedWeight>,
                             CpuGemmKeyMap>;
 
 struct MklGemmState : public Context::Resource {
@@ -665,7 +633,7 @@ void mkl_ral_gemm(ExecutionContext* ctx, void* stream_handle,
       unique_name, []() { return new MklGemmState; });
   opaque_t packed_weight;
   {
-    CpuGemmKey key{m, n, k, 1, tp_a, tp_b, B.data, kDiscCpuDefaultThreadId};
+    GEMMParamsKey key{m, n, k, 1, tp_a, tp_b, B.data, kDiscCpuDefaultThreadId};
     std::lock_guard<std::mutex> l(state->mu);
     auto& cache = state->cache;
     auto it = cache.find(key);
@@ -691,7 +659,7 @@ struct OnednnGemmState : public Context::Resource {
 
 using MatmulPrimitive = ideep::matmul_forward::super;
 using OnednnAclGemmCache =
-    ideep::utils::lru_cache<CpuGemmKey, std::shared_ptr<MatmulPrimitive>,
+    ideep::utils::lru_cache<GEMMParamsKey, std::shared_ptr<MatmulPrimitive>,
                             CpuGemmKeyMap>;
 
 struct OnednnAclGemmState : public Context::Resource {
@@ -700,7 +668,7 @@ struct OnednnAclGemmState : public Context::Resource {
 };
 
 std::shared_ptr<MatmulPrimitive> getOrCreateMatmulPrimitive(
-    OnednnAclGemmCache& cached_primitive, const CpuGemmKey& key,
+    OnednnAclGemmCache& cached_primitive, const GEMMParamsKey& key,
     const tensor& src, const tensor& weight, tensor& output) {
   auto it = cached_primitive.find(key);
   if (it == cached_primitive.end()) {
@@ -755,7 +723,8 @@ void onednn_ral_gemm(ExecutionContext* ctx, void* stream_handle,
       unique_name, []() { return new OnednnAclGemmState; });
   std::shared_ptr<MatmulPrimitive> primitive;
   {
-    CpuGemmKey key{m, n, k, 1, tp_a, tp_b, B.data, std::this_thread::get_id()};
+    GEMMParamsKey key{m,    n,    k,      1,
+                      tp_a, tp_b, B.data, std::this_thread::get_id()};
     std::lock_guard<std::mutex> l(state->mu);
     primitive = getOrCreateMatmulPrimitive(state->cached_primitive, key, src,
                                            weight, output);
@@ -929,7 +898,8 @@ void onednn_ral_batch_gemm(ExecutionContext* ctx, void* stream_handle,
   std::shared_ptr<MatmulPrimitive> primitive;
   {
     std::lock_guard<std::mutex> l(state->mu);
-    CpuGemmKey key{m, n, k, b, tp_a, tp_b, B.data, std::this_thread::get_id()};
+    GEMMParamsKey key{m,    n,    k,      b,
+                      tp_a, tp_b, B.data, std::this_thread::get_id()};
     primitive = getOrCreateMatmulPrimitive(state->cached_primitive, key, src,
                                            weight, output);
   }
diff --git a/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.h b/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.h
index d4873f834c7..ab6fcd9523e 100644
--- a/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.h
+++ b/tao_compiler/mlir/xla/ral/context/common_context_impl_mkldnn.h
@@ -84,6 +84,66 @@ inline void hash_combine(std::size_t& seed, const T& v) {
 
 extern const std::thread::id kDiscCpuDefaultThreadId;
 
+struct GEMMParamsKeyHasher;
+
+struct GEMMParamsKey {
+  int m = -1;
+  int n = -1;
+  int k = -1;
+  int batch = 1;
+  bool transpose_a = false;
+  bool transpose_b = false;
+  opaque_t const_weight_ptr = nullptr;
+  // We need this since the matmul primitive is not thead safe.
+  // To enable large parallelism, we cache the primitive per-thread.
+  std::thread::id tid;
+
+  bool operator==(const GEMMParamsKey& rhs) const {
+    return m == rhs.m && n == rhs.n && k == rhs.k && batch == rhs.batch &&
+           transpose_a == rhs.transpose_a && transpose_b == rhs.transpose_b &&
+           const_weight_ptr == rhs.const_weight_ptr && tid == rhs.tid;
+  }
+
+  using Hasher = GEMMParamsKeyHasher;
+};
+
+struct GEMMParamsKeyHasher {
+  std::size_t operator()(const GEMMParamsKey& key) const {
+    std::size_t seed = std::hash<int>()(key.m);
+    hash_combine(seed, key.n);
+    hash_combine(seed, key.k);
+    hash_combine(seed, key.batch);
+    hash_combine(seed, key.transpose_a);
+    hash_combine(seed, key.transpose_b);
+    hash_combine(seed, key.const_weight_ptr);
+    hash_combine(seed, key.tid);
+    return seed;
+  }
+};
+
+template <typename Tinput, int NDims, typename Tfilter = Tinput,
+          typename Toutput = Tinput>
+inline GEMMParamsKey makeGEMMParamsKey(MemRefType<Tinput, NDims> input,
+                                       MemRefType<Tfilter, NDims> weight,
+                                       MemRefType<Toutput, NDims> result,
+                                       bool tp_a, bool tp_b,
+                                       bool weight_is_const,
+                                       const std::thread::id& tid) {
+  GEMMParamsKey key;
+  key.tid = tid;
+  key.const_weight_ptr = weight_is_const ? weight.data : nullptr;
+  for (int i = 0; i < NDims; ++i) {
+    key.batch *= input.sizes[i];
+  }
+  key.m = tp_a ? input.sizes[NDims - 1] : input.sizes[NDims - 2];
+  key.n = tp_b ? weight.sizes[NDims - 2] : weight.sizes[NDims - 1];
+  key.k = tp_a ? input.sizes[NDims - 2] : input.sizes[NDims - 1];
+  key.transpose_a = tp_a;
+  key.transpose_b = tp_b;
+
+  return key;
+}
+
 struct ConvParams {
   format_tag input_format;
   format_tag filter_format;
@@ -102,6 +162,8 @@ struct ConvParams {
   bool is_depthwise = false;
 };
 
+struct ConvParamsKeyHasher;
+
 struct ConvParamsKey {
   std::vector<int64_t> src_dims;
   std::vector<int64_t> weight_dims;
@@ -112,6 +174,8 @@ struct ConvParamsKey {
   // We need this in case the cased kernel is not thead safe.
   // To enable large parallelism, we cache the primitive per-thread.
   std::thread::id tid;
+
+  using Hasher = ConvParamsKeyHasher;
 };
 
 inline bool operator==(const ConvParamsKey& lhs, const ConvParamsKey& rhs) {
@@ -463,8 +527,12 @@ void dumpConvLikeKernelProflingInfo(const ConvParams& params, size_t nanosec,
 }
 
 #if defined(TAO_AARCH64)
+
+// Sets thread pool configurations in the first step.
+bool applyACLThreadPoolConfigIfNotSet();
+
 template <typename OpTy>
-struct AclConvLikeInfo {
+struct AclOpInfo {
   arm_compute::Tensor src;
   arm_compute::Tensor weights;
   arm_compute::Tensor dst;
@@ -472,11 +540,11 @@ struct AclConvLikeInfo {
 };
 
 template <typename OpTy>
-class AclConvLikeThreadSafeInfo {
+class AclOpThreadSafeInfo {
  public:
-  inline std::shared_ptr<AclConvLikeInfo<OpTy>> getOrCreate(
-      std::function<std::shared_ptr<AclConvLikeInfo<OpTy>>(
-          const arm_compute::ITensorPack*)>
+  inline std::shared_ptr<AclOpInfo<OpTy>> getOrCreate(
+      std::function<
+          std::shared_ptr<AclOpInfo<OpTy>>(const arm_compute::ITensorPack*)>
           creator) {
     auto tid = std::this_thread::get_id();
     {
@@ -499,30 +567,29 @@ class AclConvLikeThreadSafeInfo {
  private:
   std::mutex mu;
   // `mainInfo` is the one owns the pre-packed weights.
-  std::shared_ptr<AclConvLikeInfo<OpTy>> mainInfo_;
-  std::unordered_map<std::thread::id, std::shared_ptr<AclConvLikeInfo<OpTy>>>
+  std::shared_ptr<AclOpInfo<OpTy>> mainInfo_;
+  std::unordered_map<std::thread::id, std::shared_ptr<AclOpInfo<OpTy>>>
       infoMap_;
 };
 
 template <typename TKey, typename TValue>
-using AclConvLikeMap = std::unordered_map<TKey, TValue, ConvParamsKeyHasher>;
-template <typename OpTy>
-using AclConvLikeCache = ideep::utils::lru_cache<
-    ConvParamsKey, std::shared_ptr<AclConvLikeInfo<OpTy>>, AclConvLikeMap>;
-template <typename OpTy>
-using AclConvLikeThreadSafeCache =
-    ideep::utils::lru_cache<ConvParamsKey,
-                            std::shared_ptr<AclConvLikeThreadSafeInfo<OpTy>>,
-                            AclConvLikeMap>;
-
-template <typename OpTy>
-struct AclConvLikeState : public Context::Resource {
+using AclOpMap = std::unordered_map<TKey, TValue, typename TKey::Hasher>;
+template <typename TKey, typename OpTy>
+using AclOpCache =
+    ideep::utils::lru_cache<TKey, std::shared_ptr<AclOpInfo<OpTy>>, AclOpMap>;
+template <typename TKey, typename OpTy>
+using AclOpThreadSafeCache =
+    ideep::utils::lru_cache<TKey, std::shared_ptr<AclOpThreadSafeInfo<OpTy>>,
+                            AclOpMap>;
+
+template <typename TKey, typename OpTy>
+struct AclOpState : public Context::Resource {
   std::mutex mu;
-  AclConvLikeCache<OpTy> cache{getWeightPrePackingCacheCapacity()};
+  AclOpCache<TKey, OpTy> cache{getWeightPrePackingCacheCapacity()};
 
-  inline std::shared_ptr<AclConvLikeInfo<OpTy>> getOrCreate(
-      const ConvParamsKey& key,
-      std::function<std::shared_ptr<AclConvLikeInfo<OpTy>>()> creator) {
+  inline std::shared_ptr<AclOpInfo<OpTy>> getOrCreate(
+      const TKey& key,
+      std::function<std::shared_ptr<AclOpInfo<OpTy>>()> creator) {
     std::lock_guard<std::mutex> l(mu);
     auto it = cache.find(key);
     if (it == cache.end()) {
@@ -532,17 +599,16 @@ struct AclConvLikeState : public Context::Resource {
   }
 };
 
-template <typename OpTy>
-struct AclConvLikeThreadSafeState : public Context::Resource {
+template <typename TKey, typename OpTy>
+struct AclOpThreadSafeState : public Context::Resource {
   std::mutex mu;
-  AclConvLikeThreadSafeCache<OpTy> cache{getWeightPrePackingCacheCapacity()};
+  AclOpThreadSafeCache<TKey, OpTy> cache{getWeightPrePackingCacheCapacity()};
 
-  inline std::shared_ptr<AclConvLikeInfo<OpTy>> getOrCreate(
-      const ConvParamsKey& key,
-      std::function<std::shared_ptr<AclConvLikeInfo<OpTy>>(
-          const arm_compute::ITensorPack*)>
-          creator) {
-    using ThreadSafeInfo = AclConvLikeThreadSafeInfo<OpTy>;
+  inline std::shared_ptr<AclOpInfo<OpTy>> getOrCreate(
+      const TKey& key, std::function<std::shared_ptr<AclOpInfo<OpTy>>(
+                           const arm_compute::ITensorPack*)>
+                           creator) {
+    using ThreadSafeInfo = AclOpThreadSafeInfo<OpTy>;
     std::shared_ptr<ThreadSafeInfo> threadSafeInfo;
     {
       std::lock_guard<std::mutex> l(mu);
@@ -560,12 +626,17 @@ struct AclConvLikeThreadSafeState : public Context::Resource {
 };
 
 using AclDepthwiseConvInfo =
-    AclConvLikeInfo<arm_compute::DISCNEDepthwiseConvolutionLayer>;
+    AclOpInfo<arm_compute::DISCNEDepthwiseConvolutionLayer>;
 using AclDepthwiseConvState =
-    AclConvLikeState<arm_compute::DISCNEDepthwiseConvolutionLayer>;
-using AclConvInfo = AclConvLikeInfo<arm_compute::DISCNEConvolutionLayer>;
+    AclOpState<ConvParamsKey, arm_compute::DISCNEDepthwiseConvolutionLayer>;
+using AclConvInfo = AclOpInfo<arm_compute::DISCNEConvolutionLayer>;
 using AclConvState =
-    AclConvLikeThreadSafeState<arm_compute::DISCNEConvolutionLayer>;
+    AclOpThreadSafeState<ConvParamsKey, arm_compute::DISCNEConvolutionLayer>;
+
+using AclQGemmInfo = AclOpInfo<arm_compute::DISCNEGEMMLowpMatrixMultiplyCore>;
+using AclQGemmState =
+    AclOpThreadSafeState<GEMMParamsKey,
+                         arm_compute::DISCNEGEMMLowpMatrixMultiplyCore>;
 
 #endif  // TAO_AARCH64
 
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 f7b2579ad45..7d5ad9e98fd 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
@@ -14,6 +14,7 @@
 #include <sstream>
 #include <thread>
 
+#include "arm_compute/core/utils/quantization/AsymmHelpers.h"
 #include "tensorflow/compiler/mlir/xla/ral/context/common_context_impl_mkldnn.h"
 
 namespace tao {
@@ -180,6 +181,8 @@ void ral_qconv_acl_s8_s8_s8_per_channel(
     ctx->signalError(Context::FAILURE, "invalid conv params");
   }
 
+  applyACLThreadPoolConfigIfNotSet();
+
   if (TAO_VLOG_IS_ON(2)) {
     for (int i = 0; i < Size(input); ++i) {
       TAO_VLOG(0) << "input[" << i
@@ -329,6 +332,139 @@ void ral_qconv_acl_s8_s8_s8_per_channel(
   }
 }
 
+void ral_qgemm_acl_s8_s8_s8_per_channel(
+    ExecutionContext* ctx, opaque_t /*stream_handle*/,
+    MemRefType<int8_t, 2> input, MemRefType<int8_t, 2> weight,
+    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, 2> result, bool tp_a, bool tp_b, bool weight_is_const) {
+  CpuTimer timer("ral_cpu_qgemm");
+  if (isEmptyMemref(input) || isEmptyMemref(weight) || isEmptyMemref(result)) {
+    TAO_VLOG(1) << "ral_cpu_batch_gemm: early return for empty tensor";
+    return;
+  }
+
+  // ACL GEMM kernel does not support transpose.
+  // TODO(disc): use standalone ACL transpose kernel to imitate.
+  if (tp_a || tp_b) {
+    ctx->signalError(Context::FAILURE,
+                     "not supported ral_qgemm with transpose");
+  }
+
+  // Make sure thread pool configuration is set.
+  applyACLThreadPoolConfigIfNotSet();
+
+  int64_t m = tp_a ? input.sizes[1] : input.sizes[0];
+  int64_t k = tp_a ? input.sizes[0] : input.sizes[1];
+  if (k != (tp_b ? weight.sizes[1] : weight.sizes[0])) {
+    ctx->signalError(Context::FAILURE, "mismatch contraction dim for gemm");
+    return;
+  }
+  int64_t n = (tp_b ? weight.sizes[0] : weight.sizes[1]);
+
+  auto AclQGemmCreator = [&](const arm_compute::ITensorPack* pack) {
+    std::shared_ptr<AclQGemmInfo> info(new AclQGemmInfo);
+    auto src_shape = TensorShape(k, m);
+    auto weights_shape = TensorShape(n, k);
+    auto dst_shape = TensorShape(n, m);
+
+    DataType data_type = DataType::QASYMM8_SIGNED;
+    TensorInfo src_info = TensorInfo(src_shape, 1, data_type);
+    TensorInfo weights_info =
+        TensorInfo(weights_shape, 1, DataType::QSYMM8_PER_CHANNEL);
+    TensorInfo dst_info = TensorInfo(dst_shape, 1, data_type);
+
+    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);
+
+    arm_compute::GEMMLowpOutputStageInfo osInfo;
+    osInfo.type =
+        arm_compute::GEMMLowpOutputStageType::QUANTIZE_DOWN_FIXEDPOINT;
+    osInfo.gemmlowp_offset = *resultZeroPoints.data;
+    osInfo.output_data_type = data_type;
+    osInfo.is_quantized_per_channel = true;
+    arm_compute::quantization::calculate_quantized_multipliers(
+        src_info.quantization_info(), weights_info.quantization_info(),
+        dst_info.quantization_info(), osInfo);
+    arm_compute::GEMMInfo gemmInfo;
+    gemmInfo.set_gemmlowp_output_stage(osInfo);
+
+    auto status = info->op.validate(&src_info, &weights_info, nullptr,
+                                    &dst_info, gemmInfo);
+    if (!status) {
+      ctx->signalError(
+          Context::FAILURE,
+          "fail to validate ral_qgemm_acl_s8_s8_s8_per_channel conv: " +
+              status.error_description());
+    } else {
+      info->op.configure(&info->src, &info->weights, nullptr, &info->dst,
+                         gemmInfo);
+    }
+    if (pack) info->op.reuse_packed_weight(*pack);
+    info->op.prepare(&info->src, &info->weights, nullptr, &info->dst);
+    return info;
+  };
+
+  std::shared_ptr<AclQGemmInfo> info;
+  if (isWeightPrePackingEnabled() && weight_is_const) {
+    std::string unique_name = "disc.ral_qgemm_acl_s8_s8_s8_per_channel";
+    auto state = ctx->getOrCreateResource<AclQGemmState>(
+        unique_name, []() { return new AclQGemmState; });
+    auto key = makeGEMMParamsKey(input, weight, result, tp_a, tp_b,
+                                 weight_is_const, kDiscCpuDefaultThreadId);
+    info = state->getOrCreate(key, AclQGemmCreator);
+  } else {
+    info = AclQGemmCreator(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);
+
+  timer.Stop();
+
+  if (isProfilingEnabled()) {
+    int64_t bytes = sizeof(int8_t) * m * k + sizeof(int8_t) * k * n +
+                    sizeof(int8_t) * m * n;
+    TAO_VLOG(0) << "ral_cpu_gemm:\n"
+                << "\tpa = " << input.data << "\n"
+                << "\tpb = " << weight.data << "\n"
+                << "\tpc = " << result.data << "\n"
+                << "\tm = " << m << "\n"
+                << "\tn = " << n << "\n"
+                << "\tk = " << k << "\n"
+                << "\ttp_a = " << tp_a << "\n"
+                << "\ttp_b = " << tp_b << "\n"
+                << "\tweight_is_const = " << weight_is_const << "\n"
+                << "\tMath Ops = " << 2 * m * n * k << "\n"
+                << "\tBytes = " << bytes << "\n"
+                << "\tBandwidth = "
+                << double(bytes) / double(timer.GetNanoSeconds()) << " GB\n"
+                << "\tGFLOPS = "
+                << double(2 * m * n * k) / double(timer.GetNanoSeconds())
+                << "\n";
+  }
+}
+
 }  // namespace
 
 // Deprecated, remove such implementation once refactor is done.
@@ -337,6 +473,8 @@ 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>);
 
+TAO_RAL_API("ral_qgemm", "cpu", ral_qgemm_acl_s8_s8_s8_per_channel);
+
 }  // namespace ral
 }  // namespace tao
 #endif


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_p_i8_per_channel.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_p_i8_per_channel.mlir
new file mode 100644
index 00000000000..d3749a0afb4
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_p_i8_per_channel.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<?x71xf32>) 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?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x?xf32>
+
+      %weight:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.8> : tensor<25x71xf32>} : () -> tensor<25x71xf32>
+      %weight_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.568627e-2> : tensor<71xf32>} : () -> tensor<71xf32>
+      %weight_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<71xi32>} : () -> tensor<71xi32>
+      %quantized_weight:2 = tf_executor.island wraps "tf.DiscFakeQuant"(%weight, %weight_scale, %weight_zero_point) {
+          axis = [1], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 127 : i64, quant_min = -128 : i64} :
+      (tensor<25x71xf32>, tensor<71xf32>, tensor<71xi32>) -> tensor<25x71xf32>
+
+      %result:2 = tf_executor.island wraps "tf.MatMul"(%quantized_input, %quantized_weight)
+      {
+        transpose_a = false,
+        transpose_b = false
+      } : (tensor<?x?xf32>, tensor<25x71xf32>) -> tensor<?x71xf32>
+
+      %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<?x71xf32>, tensor<f32>, tensor<i32>) -> tensor<?x71xf32>
+
+      tf_executor.fetch %quantized_result : tensor<?x71xf32>
+    }
+    return %graph : tensor<?x71xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_s_i8_per_channel.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_s_i8_per_channel.mlir
new file mode 100644
index 00000000000..ad9f9bc63e6
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/quantized_matmul_s_i8_per_channel.mlir
@@ -0,0 +1,33 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%input: tensor<4x25xf32>) -> (tensor<4x71xf32>) 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<4x25xf32>, tensor<f32>, tensor<i32>) -> tensor<4x25xf32>
+
+      %weight:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.8> : tensor<25x71xf32>} : () -> tensor<25x71xf32>
+      %weight_scale:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.568627e-2> : tensor<71xf32>} : () -> tensor<71xf32>
+      %weight_zero_point:2 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<71xi32>} : () -> tensor<71xi32>
+      %quantized_weight:2 = tf_executor.island wraps "tf.DiscFakeQuant"(%weight, %weight_scale, %weight_zero_point) {
+          axis = [1], use_dynamic = false, use_signed = true, use_symmetric = true, num_bits = 8 : i64, quant_max = 127 : i64, quant_min = -128 : i64} :
+      (tensor<25x71xf32>, tensor<71xf32>, tensor<71xi32>) -> tensor<25x71xf32>
+
+      %result:2 = tf_executor.island wraps "tf.MatMul"(%quantized_input, %quantized_weight)
+      {
+        transpose_a = false,
+        transpose_b = false
+      } : (tensor<4x25xf32>, tensor<25x71xf32>) -> tensor<4x71xf32>
+
+      %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<4x71xf32>, tensor<f32>, tensor<i32>) -> tensor<4x71xf32>
+
+      tf_executor.fetch %quantized_result : tensor<4x71xf32>
+    }
+    return %graph : tensor<4x71xf32>
+  }
+}
\ 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 e1fb587b7b8..2fa31ecaba4 100644
--- a/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_conv2d.cc
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_conv2d.cc
@@ -22,7 +22,7 @@ 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) {
+TEST(TFQuantziedConv2d, STATIC_SHAPE_NHWC_I8_PER_CHANNEL) {
   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>();
@@ -39,7 +39,7 @@ TEST(TFQuantziedConv2d, STATIC_SHAPE_NHWC_I8_PER_TENSOR) {
       /*expect_output_vals*/ {output}));
 }
 
-TEST(TFQuantziedConv2d, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_TENSOR) {
+TEST(TFQuantziedConv2d, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_CHANNEL) {
   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>();
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_matmul.cc b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_matmul.cc
new file mode 100644
index 00000000000..28dfb2fb2d5
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/quantized_matmul.cc
@@ -0,0 +1,57 @@
+/* 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(TFQuantziedMatMul, STATIC_SHAPE_I8_PER_CHANNEL) {
+  std::vector<float> inputs(4 * 25, -0.6);
+  tensorflow::Tensor output(tensorflow::DataType::DT_FLOAT, {4, 71});
+  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_matmul_s_i8_per_channel.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"4x25xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {inputs},
+      /*expect_output_vals*/ {output}));
+}
+
+TEST(TFQuantziedMatMul, PARTIAL_DYNAMIC_SHAPE_NHWC_I8_PER_CHANNEL) {
+  std::vector<float> inputs(4 * 25, -0.6);
+  tensorflow::Tensor output(tensorflow::DataType::DT_FLOAT, {4, 71});
+  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_matmul_p_i8_per_channel.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"4x25xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {inputs},
+      /*expect_output_vals*/ {output}));
+}
+
+}  // namespace mlir_test