COMPMID-632 Assembly: Integrate gemmlowp assembly version

Integrate generic gemmlowp assembly version for u8.

Change-Id: I17ed4494c25a132b2bac581febe1544e49b4f352
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/110114
Tested-by: BSG Visual Compute Jenkins server to access repositories on http://mpd-gerrit.cambridge.arm.com <bsgcomp@arm.com>
Reviewed-by: Pablo Tello <pablo.tello@arm.com>

3 files changed, 132 insertions, 46 deletions

diff --git a/src/core/NEON/kernels/NEGEMMLowpMatrixMultiplyKernel.cpp b/src/core/NEON/kernels/NEGEMMLowpMatrixMultiplyKernel.cpp
index 2bc251e91f..208a60ce27 100644
--- a/src/core/NEON/kernels/NEGEMMLowpMatrixMultiplyKernel.cpp
+++ b/src/core/NEON/kernels/NEGEMMLowpMatrixMultiplyKernel.cpp
@@ -49,7 +49,7 @@ namespace
 {
 Error validate_arguments(const ITensorInfo *input0, const ITensorInfo *input1, const ITensorInfo *output)
 {
-    ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input0, 1, DataType::QASYMM8, DataType::S8);
+    ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input0, 1, DataType::QASYMM8, DataType::S8, DataType::U8);
     ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input0, input1);
     ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::S32);
 
diff --git a/src/core/NEON/kernels/arm64/NEGEMMLowpAArch64Kernel.cpp b/src/core/NEON/kernels/arm64/NEGEMMLowpAArch64Kernel.cpp
index b75a8ab251..db37201687 100644
--- a/src/core/NEON/kernels/arm64/NEGEMMLowpAArch64Kernel.cpp
+++ b/src/core/NEON/kernels/arm64/NEGEMMLowpAArch64Kernel.cpp
@@ -39,6 +39,7 @@ namespace arm_compute
 {
 #include "arm_compute/core/NEON/kernels/assembly/gemm_interleaved.hpp"
 #include "arm_compute/core/NEON/kernels/assembly/kernels/a64_gemm_s8_4x4.hpp"
+#include "arm_compute/core/NEON/kernels/assembly/kernels/a64_gemm_u8_4x4.hpp"
 } // namespace arm_compute
 
 #include <arm_neon.h>
@@ -50,10 +51,101 @@ namespace arm_compute
 
 namespace arm_compute
 {
+NEGEMMLowpAArch64Kernel::NEGEMMLowpAArch64Kernel()
+    : _func(nullptr)
+{
+}
+
+void gemm_interleaved_s8(const ITensor *input0, const ITensor *input1, ITensor *output, ITensor *workspace, float alpha, float beta, bool transform_0, bool transform_1, const Window &window,
+                         const ThreadInfo &info)
+{
+    const int lda = input0->info()->strides_in_bytes().y();
+    const int ldb = input1->info()->strides_in_bytes().y();
+    const int ldc = output->info()->strides_in_bytes().y() / sizeof(int32_t);
+
+    const auto in1_ptr = reinterpret_cast<const int8_t *>(input1->buffer());
+
+    const int M = std::min(output->info()->tensor_shape().y(), static_cast<size_t>(window.y().end())) - window.y().start();
+    const int N = output->info()->tensor_shape().x();
+    const int K = input0->info()->tensor_shape().x();
+
+    // Only iterate over batches
+    Window win(window);
+    win.set(0, Window::Dimension(0, 1, 1));
+    win.set(1, Window::Dimension(0, 1, 1));
+
+    Iterator in0(input0, window);
+    Iterator out(output, window);
+
+    GemmInterleaved<gemm_s8_4x4, int8_t, int32_t> gemm(&info.cpu_info, M, N, K, !transform_1, !transform_1);
+
+    constexpr size_t alignment      = 4096;
+    const size_t     offset         = (gemm.get_working_size() + alignment - 1) * info.thread_id;
+    void            *_workspace     = workspace->buffer() + offset;
+    size_t           workspace_size = workspace->info()->total_size();
+
+    if(support::cpp11::align(alignment, gemm.get_working_size(), _workspace, workspace_size) == nullptr)
+    {
+        ARM_COMPUTE_ERROR("Not enough space to align buffer!");
+    }
+
+    execute_window_loop(win, [&](const Coordinates & id)
+    {
+        gemm.execute(reinterpret_cast<const int8_t *>(in0.ptr()), lda,
+                     reinterpret_cast<const int8_t *>(in1_ptr), ldb,
+                     reinterpret_cast<int32_t *>(out.ptr()), ldc,
+                     alpha, beta, _workspace);
+    },
+    in0, out);
+}
+
+void gemm_interleaved_u8(const ITensor *input0, const ITensor *input1, ITensor *output, ITensor *workspace, float alpha, float beta, bool transform_0, bool transform_1, const Window &window,
+                         const ThreadInfo &info)
+{
+    const int lda = input0->info()->strides_in_bytes().y();
+    const int ldb = input1->info()->strides_in_bytes().y();
+    const int ldc = output->info()->strides_in_bytes().y() / sizeof(uint32_t);
+
+    const auto in1_ptr = reinterpret_cast<const uint8_t *>(input1->buffer());
+
+    const int M = std::min(output->info()->tensor_shape().y(), static_cast<size_t>(window.y().end())) - window.y().start();
+    const int N = output->info()->tensor_shape().x();
+    const int K = input0->info()->tensor_shape().x();
+
+    // Only iterate over batches
+    Window win(window);
+    win.set(0, Window::Dimension(0, 1, 1));
+    win.set(1, Window::Dimension(0, 1, 1));
+
+    Iterator in0(input0, window);
+    Iterator out(output, window);
+
+    GemmInterleaved<gemm_u8_4x4, uint8_t, uint32_t> gemm(&info.cpu_info, M, N, K, !transform_1, !transform_1);
+
+    constexpr size_t alignment      = 4096;
+    const size_t     offset         = (gemm.get_working_size() + alignment - 1) * info.thread_id;
+    void            *_workspace     = workspace->buffer() + offset;
+    size_t           workspace_size = workspace->info()->total_size();
+
+    if(support::cpp11::align(alignment, gemm.get_working_size(), _workspace, workspace_size) == nullptr)
+    {
+        ARM_COMPUTE_ERROR("Not enough space to align buffer!");
+    }
+
+    execute_window_loop(win, [&](const Coordinates & id)
+    {
+        gemm.execute(reinterpret_cast<const uint8_t *>(in0.ptr()), lda,
+                     reinterpret_cast<const uint8_t *>(in1_ptr), ldb,
+                     reinterpret_cast<uint32_t *>(out.ptr()), ldc,
+                     alpha, beta, _workspace);
+    },
+    in0, out);
+}
+
 void NEGEMMLowpAArch64Kernel::internal_configure(const ITensor *input0, const ITensor *input1, ITensor *output, ITensor *workspace, float alpha, float beta, bool transform_0, bool transform_1)
 {
-    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input0, 1, DataType::S8);
-    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::S32);
+    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input0, 1, DataType::S8, DataType::U8);
+    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::S32, DataType::U32);
     ARM_COMPUTE_ERROR_ON_MISMATCHING_DATA_TYPES(input0, input1);
 
     _input0      = input0;
@@ -65,6 +157,19 @@ void NEGEMMLowpAArch64Kernel::internal_configure(const ITensor *input0, const IT
     _transform_0 = transform_0;
     _transform_1 = transform_1;
 
+    switch(input0->info()->data_type())
+    {
+        case DataType::S8:
+            _func = &gemm_interleaved_s8;
+            break;
+        case DataType::U8:
+            _func = &gemm_interleaved_u8;
+            break;
+        default:
+            ARM_COMPUTE_ERROR("Element size not supported");
+            break;
+    }
+
     // Configure kernel window
     Window win = calculate_max_window(*output->info());
 
@@ -85,45 +190,9 @@ void NEGEMMLowpAArch64Kernel::run(const Window &window, const ThreadInfo &info)
 {
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
     ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(INEKernel::window(), window);
+    ARM_COMPUTE_ERROR_ON(_func == nullptr);
 
-    const int lda = _input0->info()->strides_in_bytes().y();
-    const int ldb = _input1->info()->strides_in_bytes().y();
-    const int ldc = _output->info()->strides_in_bytes().y() / sizeof(int32_t);
-
-    const auto in1_ptr = reinterpret_cast<const int8_t *>(_input1->buffer());
-
-    const int M = std::min(_output->info()->tensor_shape().y(), static_cast<size_t>(window.y().end())) - window.y().start();
-    const int N = _output->info()->tensor_shape().x();
-    const int K = _input0->info()->tensor_shape().x();
-
-    // Only iterate over batches
-    Window win(window);
-    win.set(0, Window::Dimension(0, 1, 1));
-    win.set(1, Window::Dimension(0, 1, 1));
-
-    Iterator in0(_input0, window);
-    Iterator out(_output, window);
-
-    GemmInterleaved<gemm_s8_4x4, int8_t, int32_t> gemm(&info.cpu_info, M, N, K, !_transform_1, !_transform_1);
-
-    constexpr size_t alignment      = 4096;
-    const size_t     offset         = (gemm.get_working_size() + alignment - 1) * info.thread_id;
-    void            *workspace      = _workspace->buffer() + offset;
-    size_t           workspace_size = _workspace->info()->total_size();
-
-    if(support::cpp11::align(alignment, gemm.get_working_size(), workspace, workspace_size) == nullptr)
-    {
-        ARM_COMPUTE_ERROR("Not enough space to align buffer!");
-    }
-
-    execute_window_loop(win, [&](const Coordinates & id)
-    {
-        gemm.execute(reinterpret_cast<const int8_t *>(in0.ptr()), lda,
-                     reinterpret_cast<const int8_t *>(in1_ptr), ldb,
-                     reinterpret_cast<int32_t *>(out.ptr()), ldc,
-                     _alpha, _beta, workspace);
-    },
-    in0, out);
+    (*_func)(_input0, _input1, _output, _workspace, _alpha, _beta, _transform_0, _transform_1, window, info);
 }
 } // namespace arm_compute
 #endif /* ARM_COMPUTE_AARCH64_V8A */
diff --git a/src/runtime/NEON/functions/NEGEMMLowpAssemblyMatrixMultiplyCore.cpp b/src/runtime/NEON/functions/NEGEMMLowpAssemblyMatrixMultiplyCore.cpp
index 708daeb265..1bf437eb5f 100644
--- a/src/runtime/NEON/functions/NEGEMMLowpAssemblyMatrixMultiplyCore.cpp
+++ b/src/runtime/NEON/functions/NEGEMMLowpAssemblyMatrixMultiplyCore.cpp
@@ -43,6 +43,7 @@ namespace arm_compute
 #include "arm_compute/core/NEON/kernels/assembly/gemm_interleaved.hpp"
 #include "arm_compute/core/NEON/kernels/assembly/kernels/a64_gemm_s8_12x8.hpp"
 #include "arm_compute/core/NEON/kernels/assembly/kernels/a64_gemm_s8_4x4.hpp"
+#include "arm_compute/core/NEON/kernels/assembly/kernels/a64_gemm_u8_4x4.hpp"
 
 } // namespace arm_compute
 
@@ -55,8 +56,8 @@ NEGEMMLowpAssemblyMatrixMultiplyCore::NEGEMMLowpAssemblyMatrixMultiplyCore(std::
 
 void NEGEMMLowpAssemblyMatrixMultiplyCore::configure(const ITensor *a, const ITensor *b, ITensor *output)
 {
-    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(a, 1, DataType::S8);
-    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::S32);
+    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(a, 1, DataType::U8, DataType::S8);
+    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::U32, DataType::S32);
     ARM_COMPUTE_ERROR_ON_MISMATCHING_DATA_TYPES(a, b);
     ARM_COMPUTE_ERROR_ON_MSG((a)->info()->dimension(0) != (b)->info()->dimension(1), "The product AB is defined only if the number of columns in A is equal to the number of rows in B");
     ARM_COMPUTE_ERROR_ON_MSG((a)->info()->dimension(1) != (output)->info()->dimension(1), "The output matrix must have the same number of rows as the matrix A");
@@ -92,9 +93,25 @@ void NEGEMMLowpAssemblyMatrixMultiplyCore::configure(const ITensor *a, const ITe
 #elif defined(ARM_COMPUTE_AARCH64_V8A)
     if(1)
     {
-        // Configure matrix multiply kernel
-        GemmInterleaved<gemm_s8_4x4, int8_t, int32_t> gemm(&ci, M, N, K, false, false);
-        _workspace.allocator()->init(TensorInfo(TensorShape{ (gemm.get_working_size() + workspace_alignment - 1) * NEScheduler::get().num_threads() }, 1, DataType::U8));
+        switch(a->info()->data_type())
+        {
+            case DataType::S8:
+            {
+                // Configure matrix multiply kernel
+                GemmInterleaved<gemm_s8_4x4, int8_t, int32_t> gemm(&ci, M, N, K, false, false);
+                _workspace.allocator()->init(TensorInfo(TensorShape{ (gemm.get_working_size() + workspace_alignment - 1) * NEScheduler::get().num_threads() }, 1, DataType::U8));
+            }
+            break;
+            case DataType::U8:
+            {
+                // Configure matrix multiply kernel
+                GemmInterleaved<gemm_u8_4x4, uint8_t, uint32_t> gemm(&ci, M, N, K, false, false);
+                _workspace.allocator()->init(TensorInfo(TensorShape{ (gemm.get_working_size() + workspace_alignment - 1) * NEScheduler::get().num_threads() }, 1, DataType::U8));
+            }
+            break;
+            default:
+                ARM_COMPUTE_ERROR("Datatype not supported");
+        }
         _memory_group.manage(&_workspace);
         // Configure matrix multiplication kernel
         auto k = arm_compute::support::cpp14::make_unique<NEGEMMLowpAArch64Kernel>();