COMPMID-1054 Update RSH's GEMM to add batch+multi support
Change-Id: Ib9d91b77f1d51976da4449fa1e6eeeffae307353
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/127876
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Pablo Tello <pablo.tello@arm.com>
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>
diff --git a/src/core/NEON/kernels/arm_gemm/gemm_native.hpp b/src/core/NEON/kernels/arm_gemm/gemm_native.hpp
index b019279..beecb76 100644
--- a/src/core/NEON/kernels/arm_gemm/gemm_native.hpp
+++ b/src/core/NEON/kernels/arm_gemm/gemm_native.hpp
@@ -28,9 +28,12 @@
#include "arm_gemm.hpp"
#include "mergeresults.hpp"
-#include "profiler.hpp"
#include "transform.hpp"
+#ifdef CYCLE_PROFILING
+#include "profiler.hpp"
+#endif
+
namespace arm_gemm
{
// Implementation of the GemmCommon abstract class.
@@ -50,6 +53,9 @@
const unsigned int _Nsize;
const unsigned int _Ksize;
+ const unsigned int _nbatches;
+ const unsigned int _nmultis;
+
Tr _beta;
const CPUInfo *const _ci;
@@ -61,8 +67,8 @@
GemmNative(GemmNative &) = delete;
GemmNative &operator=(GemmNative &) = delete;
- GemmNative(const CPUInfo *ci, const unsigned int M, const unsigned int N, const unsigned int K, const Tr beta)
- : _Msize(M), _Nsize(N), _Ksize(K), _beta(beta), _ci(ci)
+ GemmNative(const CPUInfo *ci, const unsigned int M, const unsigned int N, const unsigned int K, const unsigned int nbatches, const unsigned int nmultis, const Tr beta)
+ : _Msize(M), _Nsize(N), _Ksize(K), _nbatches(nbatches), _nmultis(nmultis), _beta(beta), _ci(ci)
{
/* For now don't do any blocking. TODO: figure out if we should. */
k_block = K;
@@ -72,29 +78,55 @@
// Window is number of out_height blocks
unsigned int get_window_size() const override
{
- return iceildiv(_Msize, strategy::out_height);
+ return iceildiv(_Msize, strategy::out_height) * _nbatches * _nmultis;
}
// Actually execute the GEMM.
void execute(unsigned int start, unsigned int end, int) override
{
+#ifdef CYCLE_PROFILING
profiler prof;
+#endif
strategy strat(_ci);
- unsigned int M_start = start * strategy::out_height;
- unsigned int M_end = std::min(end * strategy::out_height, _Msize);
+ const unsigned int window_per_batch = iceildiv(_Msize, strategy::out_height);
+ const unsigned int window_per_multi = window_per_batch * _nbatches;
+
+ const unsigned int first_multi = start / window_per_multi;
+ const unsigned int last_multi = end / window_per_multi;
+
+ const unsigned int first_batch = (start - (first_multi * window_per_multi)) / window_per_batch;
+ const unsigned int last_batch = (end - (last_multi * window_per_multi)) / window_per_batch;
+
+ const unsigned int first_row = ((start - (first_multi * window_per_multi)) % window_per_batch) * strategy::out_height;
+ const unsigned int last_row = ((end - (last_multi * window_per_multi)) % window_per_batch) * strategy::out_height;
static_assert(std::is_same<To, Toi>::value, "gemm_native: Operand types must be the same.");
static_assert(std::is_same<Tr, Tri>::value, "gemm_native: Result types must be the same.");
- for(unsigned int y0 = M_start; y0 < M_end; y0 += strategy::out_height)
+ for(unsigned int multi = first_multi; multi <= last_multi; multi++)
{
- unsigned int ymax = std::min(y0 + strategy::out_height, M_end);
+ const unsigned int batch_0 = (multi == first_multi) ? first_batch : 0;
+ const unsigned int batch_max = (multi == last_multi) ? last_batch : _nbatches;
- prof(PROFILE_KERNEL, (ymax - y0) * _Nsize * _Ksize, [&](void)
+ for(unsigned int batch = batch_0; batch < batch_max; batch++)
{
- strat.kernel(this->_Aptr + (y0 * this->_lda), this->_lda, this->_Bptr, this->_ldb, this->_Cptr + (y0 * this->_ldc), this->_ldc, _beta, (ymax - y0), _Nsize, _Ksize);
- });
+ const unsigned int m_start = ((multi == first_multi) && (batch == first_batch)) ? first_row : 0;
+ const unsigned int m_end = ((multi == last_multi) && (batch == last_batch)) ? last_row : _Msize;
+
+ for(unsigned int y0 = m_start; y0 < m_end; y0 += strategy::out_height)
+ {
+ const unsigned int ymax = std::min(y0 + strategy::out_height, m_end);
+#ifdef CYCLE_PROFILING
+ auto p = prof.ScopedProfiler(PROFILE_KERNEL, (ymax - y0) * _Nsize * _Ksize);
+#endif
+
+ strat.kernel(this->_Aptr + (multi * this->_A_multi_stride) + (batch * this->_A_batch_stride) + (y0 * this->_lda), this->_lda,
+ this->_Bptr + (multi * this->_B_multi_stride), this->_ldb,
+ this->_Cptr + (multi * this->_C_multi_stride) + (batch * this->_C_batch_stride) + (y0 * this->_ldc), this->_ldc,
+ _beta, (ymax - y0), _Nsize, _Ksize);
+ }
+ }
}
}
};