COMPMID-2178: Update GEMM assembly code.

Perform offset reduction and requantization within the assembly wrapper.

Change-Id: I5d5b3e1f6f9ef4c71805362c57f88ff199c027a3
Signed-off-by: Georgios Pinitas <georgios.pinitas@arm.com>
Reviewed-on: https://review.mlplatform.org/c/1541
Comments-Addressed: Pablo Marquez <pablo.tello@arm.com>
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

1 files changed, 91 insertions, 38 deletions

diff --git a/src/runtime/NEON/functions/NEGEMMAssemblyDispatch.cpp b/src/runtime/NEON/functions/NEGEMMAssemblyDispatch.cpp
index 2de7d2b279..2a4498b0a9 100644
--- a/src/runtime/NEON/functions/NEGEMMAssemblyDispatch.cpp
+++ b/src/runtime/NEON/functions/NEGEMMAssemblyDispatch.cpp
@@ -74,7 +74,7 @@ std::unique_ptr<IFunction> create_function_all_types(const arm_gemm::KernelDescr
 }
 
 /** Fallback in case ACL doesn't have a function */
-template <typename TypeInput, typename TypeOutput>
+template <typename TypeInput, typename TypeOutput, class OutputStage = arm_gemm::Nothing>
 class Fallback : public NEGEMMAssemblyDispatch::IFallback
 {
 public:
@@ -82,13 +82,16 @@ public:
      *
      * @param[in]  a            Input tensor containing the Matrix A.
      * @param[in]  b            Input tensor containing the Matrix B.
+     * @param[in]  c            Input tensor containing the Matrix C.
      * @param[out] d            Output tensor to store the result of matrix multiplication.
      * @param[in]  args         Matrix multiplication information.
      * @param[in]  gemm_info    GEMM meta-data
      * @param[in]  memory_group Memory group to be used by the function.
+     * @param[in]  os           Output stage meta-data.
      */
-    void configure(const ITensor *a, const ITensor *b, ITensor *d, arm_gemm::GemmArgs<TypeOutput> args,
-                   const GEMMInfo &gemm_info, MemoryGroup &memory_group);
+    void configure(const ITensor *a, const ITensor *b, const ITensor *c, ITensor *d,
+                   arm_gemm::GemmArgs<TypeOutput> args, const GEMMInfo &gemm_info,
+                   MemoryGroup &memory_group, const OutputStage &os = {});
 
     // Inherited methods overridden:
     void run() override;
@@ -118,6 +121,10 @@ private:
     {
         nullptr
     };
+    const ITensor *_c
+    {
+        nullptr
+    };
     /** Output */
     ITensor *_d{ nullptr };
     /** GEMM workspace */
@@ -130,18 +137,19 @@ private:
     GEMMInfo _gemm_info{};
 };
 
-template <typename TypeInput, typename TypeOutput>
-void Fallback<TypeInput, TypeOutput>::configure(const ITensor *a, const ITensor *b, ITensor *d, arm_gemm::GemmArgs<TypeOutput> args,
-                                                const GEMMInfo &gemm_info, MemoryGroup &memory_group)
+template <typename TypeInput, typename TypeOutput, class OutputStage>
+void Fallback<TypeInput, TypeOutput, OutputStage>::configure(const ITensor *a, const ITensor *b, const ITensor *c, ITensor *d,
+                                                             arm_gemm::GemmArgs<TypeOutput> args, const GEMMInfo &gemm_info,
+                                                             MemoryGroup &memory_group, const OutputStage &os)
 {
     arm_gemm::GemmConfig              gemm_cfg;
-    const arm_gemm::KernelDescription gemm_kernel_info = arm_gemm::get_gemm_method<TypeInput, TypeOutput>(args);
+    const arm_gemm::KernelDescription gemm_kernel_info = arm_gemm::get_gemm_method<TypeInput, TypeOutput, OutputStage>(args, os);
     if(gemm_kernel_info.method != arm_gemm::GemmMethod::GEMV_BATCHED)
     {
         gemm_cfg.filter = gemm_kernel_info.name;
         args._cfg       = &gemm_cfg;
     }
-    _gemm_kernel_asm = arm_gemm::gemm<TypeInput, TypeOutput>(args);
+    _gemm_kernel_asm = arm_gemm::gemm<TypeInput, TypeOutput, OutputStage>(args, os);
     if(_gemm_kernel_asm == nullptr)
     {
         //configuration not supported: Leave function unconfigured:
@@ -173,6 +181,7 @@ void Fallback<TypeInput, TypeOutput>::configure(const ITensor *a, const ITensor
     _optimised_kernel = std::move(acl_gemm_wrapper);
     _a                = a;
     _b                = b;
+    _c                = c;
     _d                = d;
     _gemm_info        = gemm_info;
     // Check for pre-transposed support
@@ -185,11 +194,17 @@ void Fallback<TypeInput, TypeOutput>::configure(const ITensor *a, const ITensor
     }
 }
 
-template <typename TypeInput, typename TypeOutput>
-void Fallback<TypeInput, TypeOutput>::prepare()
+template <typename TypeInput, typename TypeOutput, class OutputStage>
+void Fallback<TypeInput, TypeOutput, OutputStage>::prepare()
 {
     if(!_is_prepared)
     {
+        // Setup up matrix bias in the assembly kernel, it's just a pointer to matrix C.
+        if(_c && _c->info()->data_type() == DataType::S32)
+        {
+            _gemm_kernel_asm->set_quantized_bias(reinterpret_cast<const int32_t *>(_c->buffer() + _c->info()->offset_first_element_in_bytes()));
+        }
+
         // Pretranspose B if required
         if(_gemm_kernel_asm->B_pretranspose_required())
         {
@@ -207,8 +222,8 @@ void Fallback<TypeInput, TypeOutput>::prepare()
     }
 }
 
-template <typename TypeInput, typename TypeOutput>
-void Fallback<TypeInput, TypeOutput>::allocate_workspace(size_t workspace_size, MemoryGroup &memory_group, size_t alignment)
+template <typename TypeInput, typename TypeOutput, class OutputStage>
+void Fallback<TypeInput, TypeOutput, OutputStage>::allocate_workspace(size_t workspace_size, MemoryGroup &memory_group, size_t alignment)
 {
     ARM_COMPUTE_ERROR_ON_MSG(workspace_size == 0, "size cannot be 0");
     _workspace.allocator()->init(TensorInfo(TensorShape{ (workspace_size + alignment /* FIXME: remove alignment after COMPMID-1088 */) }, 1, DataType::S8), alignment);
@@ -216,14 +231,14 @@ void Fallback<TypeInput, TypeOutput>::allocate_workspace(size_t workspace_size,
     _workspace.allocator()->allocate();
 }
 
-template <typename TypeInput, typename TypeOutput>
-bool Fallback<TypeInput, TypeOutput>::is_configured() const
+template <typename TypeInput, typename TypeOutput, class OutputStage>
+bool Fallback<TypeInput, TypeOutput, OutputStage>::is_configured() const
 {
     return _optimised_kernel != nullptr;
 }
 
-template <typename TypeInput, typename TypeOutput>
-void Fallback<TypeInput, TypeOutput>::run()
+template <typename TypeInput, typename TypeOutput, class OutputStage>
+void Fallback<TypeInput, TypeOutput, OutputStage>::run()
 {
     const int lda = _a->info()->strides_in_bytes().y() / sizeof(TypeInput);
     int       ldb = 0;
@@ -277,10 +292,8 @@ void Fallback<TypeInput, TypeOutput>::run()
 }
 
 template <typename TypeInput, typename TypeOutput>
-void create_function_or_arm_gemm(std::unique_ptr<IFunction>                         &acl_function,
-                                 std::unique_ptr<NEGEMMAssemblyDispatch::IFallback> &arm_gemm,
-                                 MemoryGroup &memory_group, const ITensor *a, const ITensor *b,
-                                 ITensor *d, float alpha, float beta, const GEMMInfo &gemm_info,
+void create_function_or_arm_gemm(std::unique_ptr<IFunction> &acl_function, std::unique_ptr<NEGEMMAssemblyDispatch::IFallback> &arm_gemm, MemoryGroup &memory_group,
+                                 const ITensor *a, const ITensor *b, const ITensor *c, ITensor *d, float alpha, float beta, const GEMMInfo &gemm_info,
                                  std::shared_ptr<IMemoryManager> memory_manager)
 {
     INEGEMMWrapperKernel::Params p           = INEGEMMWrapperKernel::extract_parameters(a, b, d, gemm_info);
@@ -289,15 +302,51 @@ void create_function_or_arm_gemm(std::unique_ptr<IFunction>
 
     arm_gemm::GemmArgs<TypeOutput> args(&ci, p.M, p.N, p.K, p.batches, p.multis, false, false, alpha, beta, num_threads, gemm_info.pretranpose_B());
 
-    //Try to create an ACL function:
-    acl_function = create_function_all_types(arm_gemm::get_gemm_method<TypeInput, TypeOutput>(args), a, b, d, alpha, beta, gemm_info, std::move(memory_manager));
+    // Try to create an ACL function:
+    const arm_gemm::KernelDescription gemm_kernel_info = arm_gemm::get_gemm_method<TypeInput, TypeOutput>(args);
+    acl_function                                       = create_function_all_types(gemm_kernel_info, a, b, d, alpha, beta, gemm_info, std::move(memory_manager));
 
-    //If we still don't have an ACL function:
+    // If we still don't have an ACL function:
     if(acl_function == nullptr)
     {
         //Fallback onto arm_gemm function if ACL doesn't support this method.
         auto fallback = support::cpp14::make_unique<Fallback<TypeInput, TypeOutput>>();
-        fallback->configure(a, b, d, args, gemm_info, memory_group);
+        fallback->configure(a, b, c, d, args, gemm_info, memory_group);
+        arm_gemm = std::move(fallback);
+    }
+}
+
+template <typename TypeInput, typename TypeOutput>
+void create_function_or_arm_gemm_quant(std::unique_ptr<IFunction> &acl_function, std::unique_ptr<NEGEMMAssemblyDispatch::IFallback> &arm_gemm, MemoryGroup &memory_group,
+                                       const ITensor *a, const ITensor *b, const ITensor *c, ITensor *d, float alpha, float beta, const GEMMInfo &gemm_info,
+                                       std::shared_ptr<IMemoryManager> memory_manager)
+{
+    INEGEMMWrapperKernel::Params p           = INEGEMMWrapperKernel::extract_parameters(a, b, d, gemm_info);
+    const CPUInfo               &ci          = NEScheduler::get().cpu_info();
+    unsigned int                 num_threads = NEScheduler::get().num_threads();
+
+    arm_gemm::GemmArgs<TypeOutput> args(&ci, p.M, p.N, p.K, p.batches, p.multis, false, false, alpha, beta, num_threads, gemm_info.pretranpose_B());
+
+    // Configure requantization info
+    const int32_t                 a_offset = -a->info()->quantization_info().uniform().offset;
+    const int32_t                 b_offset = -b->info()->quantization_info().uniform().offset;
+    const GEMMLowpOutputStageInfo os_info  = gemm_info.gemmlowp_output_stage();
+
+    const arm_gemm::ARequantizeLayer32 gemm_requant_info(nullptr,
+                                                         a_offset, b_offset, os_info.gemmlowp_offset,
+                                                         -os_info.gemmlowp_shift, os_info.gemmlowp_multiplier,
+                                                         os_info.gemmlowp_min_bound, os_info.gemmlowp_max_bound);
+
+    // Try to create an ACL function:
+    const arm_gemm::KernelDescription gemm_kernel_info = arm_gemm::get_gemm_method<TypeInput, TypeOutput>(args, gemm_requant_info);
+    acl_function                                       = create_function_all_types(gemm_kernel_info, a, b, d, alpha, beta, gemm_info, std::move(memory_manager));
+
+    // If we still don't have an ACL function:
+    if(acl_function == nullptr)
+    {
+        // Fallback onto arm_gemm function if ACL doesn't support this method.
+        auto fallback = support::cpp14::make_unique<Fallback<TypeInput, TypeOutput, arm_gemm::ARequantizeLayer32>>();
+        fallback->configure(a, b, c, d, args, gemm_info, memory_group, gemm_requant_info);
         arm_gemm = std::move(fallback);
     }
 }
@@ -309,11 +358,10 @@ NEGEMMAssemblyDispatch::NEGEMMAssemblyDispatch(std::shared_ptr<IMemoryManager> m
 {
 }
 
-Status NEGEMMAssemblyDispatch::validate(const ITensorInfo *a, const ITensorInfo *b, const ITensorInfo *d, float alpha, float beta, const GEMMInfo &gemm_info)
+Status NEGEMMAssemblyDispatch::validate(const ITensorInfo *a, const ITensorInfo *b, const ITensorInfo *c, const ITensorInfo *d, float alpha, float beta, const GEMMInfo &gemm_info)
 {
-    ARM_COMPUTE_UNUSED(alpha);
-    ARM_COMPUTE_UNUSED(beta);
-    ARM_COMPUTE_UNUSED(gemm_info);
+    ARM_COMPUTE_UNUSED(alpha, beta, gemm_info);
+    ARM_COMPUTE_UNUSED(c);
     ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(a, b, d);
     ARM_COMPUTE_RETURN_ERROR_ON_CPU_F16_UNSUPPORTED(a);
 #ifndef __aarch64__
@@ -324,19 +372,17 @@ Status NEGEMMAssemblyDispatch::validate(const ITensorInfo *a, const ITensorInfo
     ARM_COMPUTE_RETURN_ERROR_ON_MSG(a->data_type() == DataType::F32 && d->data_type() != DataType::F32, "Only F32 output supported for F32 input");
     ARM_COMPUTE_RETURN_ERROR_ON_MSG(a->data_type() == DataType::F16 && d->data_type() != DataType::F16, "Only F16 output supported for F16 input");
     ARM_COMPUTE_RETURN_ERROR_ON_MSG(a->data_type() == DataType::U8 && d->data_type() != DataType::U32, "Only U32 output supported for U8 input");
-    ARM_COMPUTE_RETURN_ERROR_ON_MSG(a->data_type() == DataType::QASYMM8 && d->data_type() != DataType::S32 && d->data_type() != DataType::U32, "Only U32/S32 output supported for QASYMM8 input");
     ARM_COMPUTE_RETURN_ERROR_ON_MSG(a->data_type() == DataType::S8 && d->data_type() != DataType::S32, "Only S32 output supported for S8 input");
+    ARM_COMPUTE_RETURN_ERROR_ON_MSG(a->data_type() == DataType::QASYMM8 && d->data_type() != DataType::QASYMM8, "Only QASYMM8 output supported for QASYMM8 input");
     return Status{};
 }
 
-void NEGEMMAssemblyDispatch::configure(const ITensor *a, const ITensor *b, ITensor *d, float alpha, float beta, const GEMMInfo &gemm_info)
+void NEGEMMAssemblyDispatch::configure(const ITensor *a, const ITensor *b, const ITensor *c, ITensor *d, float alpha, float beta, const GEMMInfo &gemm_info)
 {
-    ARM_COMPUTE_ERROR_ON_NULLPTR(a);
-    ARM_COMPUTE_ERROR_ON_NULLPTR(b);
-    ARM_COMPUTE_ERROR_ON_NULLPTR(d);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(a, b, d);
 
     //If we don't support a combination of data types, silently return: it is the caller's responsibility to check if configure() was successful via is_configured()
-    if(!NEGEMMAssemblyDispatch::validate(a->info(), b->info(), d->info(), alpha, beta, gemm_info))
+    if(!NEGEMMAssemblyDispatch::validate(a->info(), b->info(), c != nullptr ? c->info() : nullptr, d->info(), alpha, beta, gemm_info))
     {
         return;
     }
@@ -344,20 +390,27 @@ void NEGEMMAssemblyDispatch::configure(const ITensor *a, const ITensor *b, ITens
     switch(a->info()->data_type())
     {
         case DataType::F32:
-            create_function_or_arm_gemm<float, float>(_function, _arm_gemm, _memory_group, a, b, d, alpha, beta, gemm_info, _memory_manager);
+            create_function_or_arm_gemm<float, float>(_function, _arm_gemm, _memory_group, a, b, c, d, alpha, beta, gemm_info, _memory_manager);
             break;
 #ifdef __aarch64__
         case DataType::U8:
         case DataType::QASYMM8:
-            create_function_or_arm_gemm<uint8_t, uint32_t>(_function, _arm_gemm, _memory_group, a, b, d, alpha, beta, gemm_info, _memory_manager);
+            if(d->info()->data_type() == DataType::S32)
+            {
+                create_function_or_arm_gemm<uint8_t, uint32_t>(_function, _arm_gemm, _memory_group, a, b, c, d, alpha, beta, gemm_info, _memory_manager);
+            }
+            else
+            {
+                create_function_or_arm_gemm_quant<uint8_t, uint8_t>(_function, _arm_gemm, _memory_group, a, b, c, d, alpha, beta, gemm_info, _memory_manager);
+            }
             break;
         case DataType::S8:
-            create_function_or_arm_gemm<int8_t, int32_t>(_function, _arm_gemm, _memory_group, a, b, d, alpha, beta, gemm_info, _memory_manager);
+            create_function_or_arm_gemm<int8_t, int32_t>(_function, _arm_gemm, _memory_group, a, b, c, d, alpha, beta, gemm_info, _memory_manager);
             break;
 #endif /* __aarch64__ */
 #ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
         case DataType::F16:
-            create_function_or_arm_gemm<float16_t, float16_t>(_function, _arm_gemm, _memory_group, a, b, d, alpha, beta, gemm_info, _memory_manager);
+            create_function_or_arm_gemm<float16_t, float16_t>(_function, _arm_gemm, _memory_group, a, b, c, d, alpha, beta, gemm_info, _memory_manager);
             break;
 #endif /* __ARM_FEATURE_FP16_VECTOR_ARITHMETIC */
         default: