blob: 442739b55e4edf151fb4ac0be5a4af7fa828abea [file] [log] [blame]
Pablo Telloeb82fd22018-02-23 13:43:50 +00001/*
David Mansellaaa9da12023-03-10 13:48:50 +00002 * Copyright (c) 2017-2023 Arm Limited.
Pablo Telloeb82fd22018-02-23 13:43:50 +00003 *
4 * SPDX-License-Identifier: MIT
5 *
6 * Permission is hereby granted, free of charge, to any person obtaining a copy
7 * of this software and associated documentation files (the "Software"), to
8 * deal in the Software without restriction, including without limitation the
9 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10 * sell copies of the Software, and to permit persons to whom the Software is
11 * furnished to do so, subject to the following conditions:
12 *
13 * The above copyright notice and this permission notice shall be included in all
14 * copies or substantial portions of the Software.
15 *
16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22 * SOFTWARE.
23 */
24#pragma once
25
Pablo Telloeb82fd22018-02-23 13:43:50 +000026#include <algorithm>
David Mansell318c9f42020-07-08 13:28:45 +010027#include <cassert>
Pablo Telloeb82fd22018-02-23 13:43:50 +000028
29#include "arm_gemm.hpp"
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +000030#include "bfloat.hpp"
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000031#include "convolver.hpp"
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +000032#include "kernel_weight_format.hpp"
Viet-Hoa Do03b29712022-06-01 11:47:14 +010033#include "kernel_traits.hpp"
Pablo Telloeb82fd22018-02-23 13:43:50 +000034#include "mergeresults.hpp"
David Mansell318c9f42020-07-08 13:28:45 +010035#include "performance_parameters.hpp"
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000036#include "quantized.hpp"
Pablo Telloeb82fd22018-02-23 13:43:50 +000037#include "transform.hpp"
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000038#include "utils.hpp"
Pablo Telloeb82fd22018-02-23 13:43:50 +000039
Michalis Spyroue7e96e02018-04-13 13:44:10 +010040#ifdef CYCLE_PROFILING
41#include "profiler.hpp"
42#endif
43
Pablo Telloeb82fd22018-02-23 13:43:50 +000044// Some macros used to decide how much working space to allocate.
45// Round allocations up to the next cache line.
Anthony Barbier5f707732018-07-03 16:22:02 +010046#define ALLOC_ROUND 64
47#define ROUND_UP(x) ((((x) + ALLOC_ROUND-1) / ALLOC_ROUND) * ALLOC_ROUND)
Pablo Telloeb82fd22018-02-23 13:43:50 +000048
49// Implementation of the GemmCommon abstract class.
50//
51// This implementation interleaves the source matrices in blocks - good for
52// larger matrices.
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000053
Anthony Barbier5f707732018-07-03 16:22:02 +010054namespace arm_gemm {
55
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000056namespace {
57
58// Some kernels output to a linear buffer and require a separate merge step.
59// Others output directly to the matrix result. This helper class calls the
60// appropriate functions, using templating to avoid calling non-existent
61// functions.
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +000062template<bool MergeStep, bool FixedFormat, typename OutputStage>
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000063class kernel_and_merge {
64public:
65 template<typename strategy, typename To, typename Tr, typename Tri, typename Tab>
66 static void run (
67#ifdef CYCLE_PROFILING
68 profiler &prof,
69#endif
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +000070 strategy &strat, const To *a_ptr, const To *b_panel, size_t b_stride, Tri *c_panel,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000071 Tr *c_ptr, int ldc, int kern_k, unsigned int m_0,
72 unsigned int m_max, unsigned int n_0, unsigned int n_max, const Tr *biasptr,
73 const Activation &act, bool accumulate, const OutputStage &os, const int32_t *col_bias,
74 Tab *acc_buff);
75};
76
77// Run a kernel and call the separate merge step
78template<>
79template<typename strategy, typename To, typename Tr, typename Tri, typename Tab>
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +000080void kernel_and_merge<true, false, Nothing>::run(
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000081#ifdef CYCLE_PROFILING
82 profiler &prof,
83#endif
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +000084 strategy &strat, const To *a_ptr, const To *b_panel, size_t, Tri *c_panel,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +000085 Tr *c_ptr, int ldc, int kern_k, unsigned int m_0,
86 unsigned int m_max, unsigned int n_0, unsigned int n_max, const Tr *biasptr,
87 const Activation &act, bool accumulate, const Nothing &, const int32_t *, Tab *)
88{
89 const int bblocks = iceildiv(n_max - n_0, strategy::out_width());
90
91 {
92#ifdef CYCLE_PROFILING
93 auto p=prof.ScopedProfiler(PROFILE_KERNEL, (strategy::out_height() * bblocks * strategy::out_width() * kern_k));
94#endif
95
96 strat.kernel(a_ptr, b_panel, c_panel, 1, bblocks, kern_k);
97 }
98
99 {
100#ifdef CYCLE_PROFILING
101 auto p=prof.ScopedProfiler(PROFILE_MERGE, (strategy::out_height() * bblocks * strategy::out_width() * sizeof(Tr)));
102#endif
103 strat.transforms.Merge(c_ptr, c_panel, ldc, m_0, m_max, n_0, n_max, biasptr, act, accumulate);
104 }
105}
106
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000107// Run a fixed-format kernel and call the separate merge step
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000108template<>
109template<typename strategy, typename To, typename Tr, typename Tri, typename Tab>
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000110void kernel_and_merge<true, true, Nothing>::run(
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000111#ifdef CYCLE_PROFILING
112 profiler &prof,
113#endif
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000114 strategy &strat, const To *a_ptr, const To *b_panel, size_t b_stride, Tri *c_panel,
115 Tr *c_ptr, int ldc, int kern_k, unsigned int m_0,
116 unsigned int m_max, unsigned int n_0, unsigned int n_max, const Tr *biasptr,
117 const Activation &act, bool accumulate, const Nothing &, const int32_t *, Tab *)
118{
119 {
120#ifdef CYCLE_PROFILING
121 const int bblocks = iceildiv(n_max - n_0, strategy::out_width());
122 auto p=prof.ScopedProfiler(PROFILE_KERNEL, (strategy::out_height() * bblocks * strategy::out_width() * kern_k));
123#endif
124
125 strat.kernel(a_ptr, b_panel, b_stride, c_panel, 1, (n_max - n_0), kern_k);
126 }
127
128 {
129#ifdef CYCLE_PROFILING
130 const int bblocks = iceildiv(n_max - n_0, strategy::out_width());
131 auto p=prof.ScopedProfiler(PROFILE_MERGE, (strategy::out_height() * bblocks * strategy::out_width() * sizeof(Tr)));
132#endif
133 strat.transforms.Merge(c_ptr, c_panel, ldc, m_0, m_max, n_0, n_max, biasptr, act, accumulate);
134 }
135}
136
137// Run a kernel with integrated merge
138template<>
139template<typename strategy, typename To, typename Tr, typename Tri, typename Tab>
140void kernel_and_merge<false, false, Nothing>::run(
141#ifdef CYCLE_PROFILING
142 profiler &prof,
143#endif
144 strategy &strat, const To *a_ptr, const To *b_panel, size_t, Tri *,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000145 Tr *c_ptr, int ldc, int kern_k, unsigned int m_0, unsigned int m_max,
146 unsigned int n_0, unsigned int n_max, const Tr *biasptr,
147 const Activation &act, bool accumulate, const Nothing &, const int32_t *,
148 Tab *acc_buff)
149{
150#ifdef CYCLE_PROFILING
151 auto p=prof.ScopedProfiler(PROFILE_KERNEL, (m_max - m_0) * (n_max - n_0) * kern_k);
152#endif
153
154 // We need to offset the C pointer, but as it might be NULL (requesting output to accumulation buffer) we need
155 // to be careful not to offset a null pointer.
156 Tri *offset_c_ptr;
157
158 if (c_ptr == nullptr) {
159 offset_c_ptr = nullptr;
160 } else {
161 offset_c_ptr = c_ptr + m_0 * ldc + n_0;
162 }
163
164 strat.kernel(// A and B pointers are just the packed panels.
165 a_ptr, b_panel,
166 // Provide relevant part of output array and row stride.
167 offset_c_ptr, ldc,
168 // M, N, K sizes
169 m_max-m_0, n_max - n_0, kern_k,
170 // Bias, activation, accumulation. Need to offset the bias as needed.
171 biasptr ? biasptr + n_0 : nullptr, act, accumulate,
172 // Accumulation buffer.
173 acc_buff );
174}
175
176// Run a kernel with integrated merge, quantizing
177template<>
178template<typename strategy, typename To, typename Tr, typename Tri, typename Tab>
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000179void kernel_and_merge<false, false, Requantize32>::run(
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000180#ifdef CYCLE_PROFILING
181 profiler &prof,
182#endif
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000183 strategy &strat, const To *a_ptr, const To *b_panel, size_t, Tri *,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000184 Tr *c_ptr, int ldc, int kern_k, unsigned int m_0, unsigned int m_max,
185 unsigned int n_0, unsigned int n_max, const Tr *,
186 const Activation &, bool accumulate, const Requantize32 &qp, const int32_t *col_bias,
187 Tab *acc_buff)
188{
189#ifdef CYCLE_PROFILING
190 auto p=prof.ScopedProfiler(PROFILE_KERNEL, (m_max - m_0) * (n_max - n_0) * kern_k);
191#endif
192
193 strat.kernel(// A and B pointers are just the packed panels.
194 a_ptr, b_panel,
195 // Provide relevant part of output array and row stride.
196 c_ptr + m_0 * ldc + n_0, ldc,
197 // M, N, K sizes
198 m_max-m_0, n_max - n_0, kern_k,
199 // Bias, activation, accumulation. Need to offset the bias as needed.
200 col_bias + n_0, qp, n_0, accumulate, acc_buff);
201}
202
203// Run a kernel and call the separate quantize step
204template<>
205template<typename strategy, typename To, typename Tr, typename Tri, typename Tab>
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000206void kernel_and_merge<true, false, Requantize32>::run(
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000207#ifdef CYCLE_PROFILING
208 profiler &prof,
209#endif
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000210 strategy &strat, const To *a_ptr, const To *b_panel, size_t, Tri *c_panel,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000211 Tr *c_ptr, int ldc, int kern_k, unsigned int m_0,
212 unsigned int m_max, unsigned int n_0, unsigned int n_max, const Tr *,
213 const Activation &, bool, const Requantize32 &qp, const int32_t *col_bias,
214 Tab *)
215{
216 const int bblocks = iceildiv(n_max - n_0, strategy::out_width());
217
218 {
219#ifdef CYCLE_PROFILING
220 auto p=prof.ScopedProfiler(PROFILE_KERNEL, (strategy::out_height() * bblocks * strategy::out_width() * kern_k));
221#endif
222
223 strat.kernel(a_ptr, b_panel, c_panel, 1, bblocks, kern_k);
224 }
225
226 {
227#ifdef CYCLE_PROFILING
Georgios Pinitas4ee8b152021-07-16 16:16:43 +0100228 auto p=prof.ScopedProfiler(PROFILE_QUANTIZE, ((m_max-m_0) * bblocks * strategy::out_width() * sizeof(Tr)));
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000229#endif
230 // The interleaved kernel outputs in blocks - each block is a
231 // row-major matrix of size out_width * out_height. The merge
232 // kernels are designed to deal with this but the requantizer is
233 // not, so we need to requantize one block at a time.
234 for (int i=0; i<bblocks; i++) {
235 unsigned int n_start = n_0 + (strategy::out_width() * i);
236 unsigned int n_end = std::min(n_start + strategy::out_width(), n_max);
237
238 // The row bias is interleaved with the transposed A data, get a pointer to it here.
239 const int32_t *row_bias = reinterpret_cast<const int32_t *>(a_ptr + strategy::out_height() * kern_k);
240
241 requantize_block_32(qp, (n_end - n_start), (m_max-m_0),
242 c_panel + (i * strategy::out_width() * strategy::out_height()), strategy::out_width(),
243 c_ptr + m_0 * ldc + n_start, ldc,
244 row_bias, col_bias + n_start, n_start);
245 }
246 }
247}
248
249// Integer GEMMs can be used in two contexts - "normal" where the full 32-bit output is required, or in
250// "requantizing" context where the output will be requantized.
251//
252// These require different input transforms, as if we are requantizing we want to sum the rows of the A input, and
253// if we are not we don't.
254//
255// This helper class allows the appropriate transforms to be found, without requiring kernels that don't support
256// quantization to define useless "quantized" transforms.
257template<typename strategy, bool quantized>
258class transform_type {
259public:
260 typedef decltype(strategy::transforms) type;
261};
262
263template<typename strategy>
264class transform_type<strategy, true> {
265public:
266 typedef decltype(strategy::transforms_quantized) type;
267};
268
269// We need a similar trick here to figure out what type the accumulator buffer should be.
David Mansellaaa9da12023-03-10 13:48:50 +0000270template<typename strategy, typename OutputStage, bool ForceFloat>
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000271class accumulate_buffer_type {
272public:
273 typedef typename strategy::result_type type;
274};
275
276template<typename strategy>
David Mansellaaa9da12023-03-10 13:48:50 +0000277class accumulate_buffer_type<strategy, Requantize32, false> {
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000278public:
279 typedef int32_t type;
280};
281
David Mansellaaa9da12023-03-10 13:48:50 +0000282template<typename strategy, typename OutputStage>
283class accumulate_buffer_type<strategy, OutputStage, true> {
284public:
285 typedef float type;
286};
287
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000288// Stripe width is a concept only needed for FixedFormat kernels. Use an accessor to avoid issues in other scenarios.
289template<typename strategy, bool FixedFormat>
290struct get_stripe_width {
291 static unsigned int get() {
292 return 0;
293 }
294};
295
296template<typename strategy>
297struct get_stripe_width<strategy, true> {
298 static unsigned int get() {
299 return strategy::stripe_width();
300 }
301};
302
303// KernelWeightFormat is a similar story.
304template<typename strategy, bool FixedFormat, typename To>
305struct get_kernel_weight_format {
306 static KernelWeightFormat get() {
307 return KernelWeightFormat::NON_FIXED;
308 }
309};
310
311template<typename strategy, typename To>
312struct get_kernel_weight_format<strategy, true, To> {
313 static KernelWeightFormat get() {
314 KernelWeightFormat kwf = strategy::kernel_weight_format();
315
316 // If we are using a BF16 kernel to do an FP32 problem (fast mode) then we need to set the BF16 flag on the
317 // weight format.
318 if (std::is_same<To, float>::value && std::is_same<typename strategy::operand_type, bfloat16>::value) {
319 uint32_t kwf_i = static_cast<uint32_t>(kwf);
320 kwf_i |= 0x10;
321 kwf = static_cast<KernelWeightFormat>(kwf_i);
322 }
323
324 return kwf;
325 }
326};
327
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000328} // anonymous namespace
329
David Mansellaaa9da12023-03-10 13:48:50 +0000330template<typename strategy, typename To, typename Tr, typename OutputStage=Nothing, bool MergeStep=true, bool FixedFormat=false, bool ForceThreadColumns=false, bool ForceFloatAccumulate=false>
Anthony Barbier5f707732018-07-03 16:22:02 +0100331class GemmInterleaved : public GemmCommon<To, Tr> {
Pablo Telloeb82fd22018-02-23 13:43:50 +0000332 typedef typename strategy::operand_type Toi;
Anthony Barbier5f707732018-07-03 16:22:02 +0100333 typedef typename strategy::result_type Tri;
David Mansellaaa9da12023-03-10 13:48:50 +0000334 typedef typename accumulate_buffer_type<strategy, OutputStage, ForceFloatAccumulate>::type Tab;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000335
336 /* const properties set by constructor */
Anthony Barbier5f707732018-07-03 16:22:02 +0100337 const CPUInfo * const _ci;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000338
339 const unsigned int _Msize;
340 const unsigned int _Nsize;
341 const unsigned int _Ksize;
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000342 const unsigned int _Ksections;
343 const unsigned int _Ktotal;
344 const unsigned int _rounded_Ksize;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000345
Michalis Spyroue7e96e02018-04-13 13:44:10 +0100346 const unsigned int _nbatches;
347 const unsigned int _nmulti;
348
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000349 const bool _thread_columns;
350
Georgios Pinitas48b3ef82019-10-14 19:03:09 +0100351 const Activation _act;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000352
Anthony Barbier5f707732018-07-03 16:22:02 +0100353 const int _maxthreads;
354 int _nthreads;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000355
356 /* Blocking info */
Anthony Barbier5f707732018-07-03 16:22:02 +0100357 unsigned int _k_block=0;
358 unsigned int _x_block=0;
359 unsigned int _Mround=0;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000360
361 /* Working space, pretransposed buffer, buffer manager */
Anthony Barbier5f707732018-07-03 16:22:02 +0100362 const Toi *_B_transposed=nullptr;
Anthony Barbier5f707732018-07-03 16:22:02 +0100363 void *_working_space=nullptr;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000364
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000365 Tab *_accumulation_buffer=nullptr;
366
367 /* Output stage */
368 OutputStage _os;
369
370 /* Quantized support (in addition to 'output stage' above */
371 int32_t *col_bias = nullptr;
372
373 /* Indirect parameters. _indirect_buf doubles as a flag to indicate that "indirect" transform should be used. */
374 const To * const * const * _indirect_buf = nullptr;
375
376 /* Convolver - only set up for convolution problems, so also doubles as a flag. */
377 std::unique_ptr<convolver<To>> _convolver = nullptr;
378
379 unsigned int get_col_sum_size() const {
380 if (std::is_same<OutputStage, Requantize32>::value) {
381 return _Nsize * _nmulti * sizeof(int32_t);
382 } else {
383 return 0;
384 }
385 }
386
Pablo Telloeb82fd22018-02-23 13:43:50 +0000387 /* We will need to walk through the blocks of B in a few contexts, so
388 * factor that out. */
Anthony Barbier5f707732018-07-03 16:22:02 +0100389 class blockwalker {
Pablo Telloeb82fd22018-02-23 13:43:50 +0000390 private:
Michalis Spyroue7e96e02018-04-13 13:44:10 +0100391 /* Size loops, etc. based on our parent's configuration */
David Mansellaaa9da12023-03-10 13:48:50 +0000392 const GemmInterleaved<strategy, To, Tr, OutputStage, MergeStep, FixedFormat, ForceThreadColumns, ForceFloatAccumulate> &_parent;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000393
Anthony Barbier5f707732018-07-03 16:22:02 +0100394 /* K, X and multi parameters for current iteration. */
395 unsigned int _k0=0, _x0=0, _multi=0;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000396
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000397 /* Range of X to iterate over - used in "ForceThreadColumns" cases */
398 unsigned int _x_start=0;
399 unsigned int _x_end=_parent._Nsize;
400
Anthony Barbier5f707732018-07-03 16:22:02 +0100401 unsigned int _index=0;
402 bool _done=false;
403 bool _newkblock=true;
404 bool _newmulti=true;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000405
406 public:
David Mansellaaa9da12023-03-10 13:48:50 +0000407 blockwalker(const GemmInterleaved<strategy, To, Tr, OutputStage, MergeStep, FixedFormat, ForceThreadColumns, ForceFloatAccumulate> &parent) : _parent(parent) { }
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000408
David Mansellaaa9da12023-03-10 13:48:50 +0000409 blockwalker(const GemmInterleaved<strategy, To, Tr, OutputStage, MergeStep, FixedFormat, ForceThreadColumns, ForceFloatAccumulate> &parent,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000410 unsigned int x_start, unsigned int x_end) : _parent(parent), _x0 (_x_start), _x_start(x_start), _x_end(x_end) { }
Pablo Telloeb82fd22018-02-23 13:43:50 +0000411
Anthony Barbier5f707732018-07-03 16:22:02 +0100412 unsigned int xmax() {
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000413 return std::min(_x0 + _parent._x_block, _x_end);
Pablo Telloeb82fd22018-02-23 13:43:50 +0000414 }
415
Anthony Barbier5f707732018-07-03 16:22:02 +0100416 unsigned int kmax() {
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000417 return std::min(_k0 + _parent._k_block, _parent._Ktotal);
Pablo Telloeb82fd22018-02-23 13:43:50 +0000418 }
419
420 /* Advance to the next block, return false at the end. */
Anthony Barbier5f707732018-07-03 16:22:02 +0100421 bool advance(void) {
422 if (_done) {
Pablo Telloeb82fd22018-02-23 13:43:50 +0000423 return false;
424 }
425
Anthony Barbier5f707732018-07-03 16:22:02 +0100426 _newkblock=false;
Michalis Spyroue7e96e02018-04-13 13:44:10 +0100427 _x0 += _parent._x_block;
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000428 if (_x0 >= _x_end) {
429 _x0=_x_start;
Michalis Spyroue7e96e02018-04-13 13:44:10 +0100430 _k0 += _parent._k_block;
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000431 if (_k0 >= _parent._Ktotal) {
Anthony Barbier5f707732018-07-03 16:22:02 +0100432 _k0=0;
Michalis Spyroue7e96e02018-04-13 13:44:10 +0100433 _multi++;
Anthony Barbier5f707732018-07-03 16:22:02 +0100434 if (_multi >= _parent._nmulti) {
435 _done=true;
Michalis Spyroue7e96e02018-04-13 13:44:10 +0100436 return false;
437 }
Anthony Barbier5f707732018-07-03 16:22:02 +0100438 _newmulti=true;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000439 }
Anthony Barbier5f707732018-07-03 16:22:02 +0100440 _newkblock=true;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000441 }
442 _index++;
443
444 return true;
445 }
446
Anthony Barbier5f707732018-07-03 16:22:02 +0100447 unsigned int k0(void) { return _k0; }
448 unsigned int x0(void) { return _x0; }
449 unsigned int multi(void) { return _multi; }
450 unsigned int index(void) { return _index; }
451 bool done(void) { return _done; }
452 bool newkblock(void) { return _newkblock; }
Pablo Telloeb82fd22018-02-23 13:43:50 +0000453 };
454
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000455 // "k block" has two distinct uses: figuring out which iterations of K
456 // to actually process, but also various size/pointer computations. The
457 // latter needs to take account of the extra space needed for the row
458 // sums, if appropriate.
459 unsigned int get_total_k_depth() const {
460 unsigned int k_depth = _k_block;
461
462 if (std::is_same<OutputStage, Requantize32>::value) {
463 k_depth += sizeof(int32_t) / sizeof(Toi);
464 }
465
466 return k_depth;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000467 }
468
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000469 // A working size.
470 size_t get_a_working_size() const {
471 if (_thread_columns) {
472 // For 2D threading: allocate a buffer of one block of rows per thread
473 return ROUND_UP(sizeof(Toi) * get_total_k_depth() * strategy::out_height() * _maxthreads);
474 } else {
475 // For 1D threaded: one of these needed, regardless of thread count. Divided according to window.
476 return ROUND_UP(sizeof(Toi) * get_total_k_depth() * _Mround * _nbatches);
477 }
478 }
479
480 // C working size: One needed per thread. Not needed if there is no merge step.
Anthony Barbier5f707732018-07-03 16:22:02 +0100481 size_t get_c_working_size() const {
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000482 if (MergeStep) {
483 return ROUND_UP(sizeof(Tri) * _x_block * strategy::out_height());
484 } else {
485 return 0;
486 }
487 }
488
489 // Accumulation buffer size
490 size_t get_accumulation_buffer_size() const {
491 // We only support an accumulation buffer for non-merge cases.
492 if (MergeStep) {
493 return 0;
494 }
495
496 // Check if we are actually blocking
497 if (_k_block == _Ktotal) {
498 return 0;
499 }
500
501 // We are no-merge, non-quantized with active blocking: accumulation buffer needed.
502 size_t size_per_buffer = sizeof(Tab) * strategy::out_height() * strategy::out_width();
503 size_t num_buffers = iceildiv(_Msize, strategy::out_height()) * iceildiv(_Nsize, strategy::out_width()) * _nbatches * _nmulti;
504
505 return num_buffers * size_per_buffer;
506 }
507
508 // Get pointer into accumulation buffer
509 Tab *get_accumulation_buffer(unsigned int M, unsigned int N, unsigned int batch, unsigned int multi) const {
510 // Don't do anything if there's no buffer.
511 if (_accumulation_buffer == nullptr) {
512 return nullptr;
513 }
514
515 // Here we are indexing an appropriately sized pointer, so no sizeof() needed to convert to bytes.
516 size_t size_per_buffer = strategy::out_height() * strategy::out_width();
517
518 size_t buffer_rows = iceildiv(_Msize, strategy::out_height());
519 size_t buffer_cols = iceildiv(_Nsize, strategy::out_width());
520 size_t buffers_per_batch = (buffer_rows * buffer_cols);
521 size_t buffers_per_multi = buffers_per_batch * _nbatches;
522
523 // M/N must reference the top-left corner of a block.
524 size_t row = M / strategy::out_height();
525 assert(M % strategy::out_height() == 0);
526 size_t col = N / strategy::out_width();
527 assert(N % strategy::out_width() == 0);
528
529 size_t buffer_index = multi * buffers_per_multi + batch * buffers_per_batch + row * buffer_cols + col;
530
531 return _accumulation_buffer + (buffer_index * size_per_buffer);
532 }
533
534 int32_t row_sum_multiplier() const {
535 if (std::is_same<OutputStage, Requantize32>::value) {
536 const Requantize32 *qp = reinterpret_cast<const Requantize32 *>(&_os);
537
538 return -qp->b_offset;
539 }
540
541 return 0;
542 }
543
544 // Heuristics to decide whether to use the 'thread columns' regime
545 static bool is_thread_columns(const GemmArgs &args) {
546 // For now, there is a templace parameter to force it.
547 if (ForceThreadColumns) {
548 return true;
549 }
550
551 // Never do this for single threaded cases.
552 if (args._maxthreads == 1) {
553 return false;
554 }
555
556 // How many blocks of work are available for threading on M?
557 int m_blocks = iceildiv(args._Msize, strategy::out_height()) * args._nbatches;
558
559 // If we just can't share the work across threads with the row threading regime.
560 if (args._maxthreads > m_blocks) {
561 return true;
562 }
563
564 // If the row threading regime is too wasteful (20% threshold)
565 if (((roundup(m_blocks, args._maxthreads) * 100) / m_blocks) > 120) {
566 return true;
567 }
568
569 return false;
570 }
571
572 static unsigned int get_ktotal(const GemmArgs &args) {
573 return args._Ksections * roundup(args._Ksize, strategy::k_unroll());
Pablo Telloeb82fd22018-02-23 13:43:50 +0000574 }
575
David Mansell318c9f42020-07-08 13:28:45 +0100576 static unsigned int get_k_block_size(const GemmArgs &args) {
577 if (args._cfg && args._cfg->inner_block_size) {
Georgios Pinitas4ee8b152021-07-16 16:16:43 +0100578 return roundup(args._cfg->inner_block_size, strategy::k_unroll());
David Mansell318c9f42020-07-08 13:28:45 +0100579 }
580
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000581 // K blocking not supported if we are requantizing.
582 if (std::is_same<OutputStage, Requantize32>::value) {
583 return get_ktotal(args);
584 }
585
Viet-Hoa Do03b29712022-06-01 11:47:14 +0100586 // Special blocking for SME
587 if (is_sme<strategy>::value) {
588 // Don't bother to block below this size threshold, experimentally determined to be 320 for FP32
589 unsigned int scaling_threshold = 1280 / sizeof(Toi);
590
591 if (get_ktotal(args) <= scaling_threshold) {
592 return get_ktotal(args);
593 }
594
595 // Once we are blocking, this (lower) threshold determines when we should use more blocks
596 // NOTE: Could be that some factor-based solution would work better here.
597 unsigned int max_block_size = 1024 / sizeof(Toi);
598
599 unsigned int num_k_blocks = iceildiv(get_ktotal(args), max_block_size);
600
601 unsigned int k_block = roundup(iceildiv(get_ktotal(args), num_k_blocks), strategy::k_unroll());
602
603 return k_block;
604 }
605
David Mansell318c9f42020-07-08 13:28:45 +0100606 const unsigned int L1_size = args._ci->get_L1_cache_size();
607 unsigned int k_block;
608
609 // k_block: Find out how much of the larger array can be loaded into half the cache.
610 // This should account for associative caches.
611 k_block = (L1_size / 2) / (sizeof(Toi) * (std::max(strategy::out_width(), strategy::out_height())));
612
613 // Needs to be (at least a single) multiple of the K unroll level.
614 k_block /= strategy::k_unroll();
615 k_block = std::max(k_block, 1U) * strategy::k_unroll();
616
617 // Now tune to presented problem size; this is how many blocks we need.
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000618 unsigned int num_k_blocks = iceildiv(get_ktotal(args), k_block);
David Mansell318c9f42020-07-08 13:28:45 +0100619
620 // So divide the space equally into that many blocks.
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000621 k_block = iceildiv(get_ktotal(args), num_k_blocks);
David Mansell318c9f42020-07-08 13:28:45 +0100622
623 // And round UP to the K unroll level required.
624 k_block = roundup(k_block, strategy::k_unroll());
625
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000626 assert(k_block > 0);
627
David Mansell318c9f42020-07-08 13:28:45 +0100628 return k_block;
629 }
Pablo Telloeb82fd22018-02-23 13:43:50 +0000630
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000631 static unsigned int get_x_block_size(const GemmArgs &args) {
632 if (is_thread_columns(args)) {
633 // In 2D mode, override X block, because we will process width first.
634 return roundup(args._Nsize, strategy::out_width());
635 }
636
637 if (args._cfg && args._cfg->outer_block_size) {
638 return roundup(args._cfg->outer_block_size, strategy::out_width());
639 }
640
641 unsigned int x_block;
642 const unsigned int L2_size = args._ci->get_L2_cache_size();
643 const unsigned int k_block = get_k_block_size(args);
644
645 // x_block: Work out how many rows (of length k_block) will fit in the L2
646 // Don't allocate more than 90% of the L2 to allow for overheads, and subtract off the L1 contents.
647 const unsigned int scaled_l2_size = (L2_size * 9) / 10;
648 const unsigned int k_block_area = k_block * sizeof(Toi) * (strategy::out_width() + strategy::out_height());
649
650 // .. if the L1 contents is bigger than the L2, just return a minimal size block.
651 if (k_block_area > scaled_l2_size) {
652 return strategy::out_width();
653 }
654
655 x_block = (scaled_l2_size - k_block_area) / (sizeof(Toi) * k_block);
656
657 // Needs to be (at least a single) multiple of the kernel output width.
658 x_block /= strategy::out_width();
659 x_block = std::max(x_block, 1u) * strategy::out_width();
660
661 // And tune to the presented problem size.
662 unsigned int num_x_blocks = iceildiv(args._Nsize, x_block);
663 x_block = iceildiv(args._Nsize, num_x_blocks);
664
665 x_block = roundup(x_block, strategy::out_width());
666
667 assert(x_block > 0);
668
669 return x_block;
670 }
671
Pablo Telloeb82fd22018-02-23 13:43:50 +0000672public:
673 GemmInterleaved(GemmInterleaved &) = delete;
Anthony Barbier5f707732018-07-03 16:22:02 +0100674 GemmInterleaved & operator= (GemmInterleaved &) = delete;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000675
676 /* Constructor */
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000677 GemmInterleaved(const GemmArgs &args, const OutputStage &os)
678 : _ci(args._ci), _Msize(args._Msize), _Nsize(args._Nsize), _Ksize(args._Ksize),
679 _Ksections(args._Ksections), _Ktotal(get_ktotal(args)),
680 _rounded_Ksize(roundup(_Ksize, strategy::k_unroll())),
681 _nbatches(args._nbatches), _nmulti(args._nmulti), _thread_columns(is_thread_columns(args)),
682 _act(args._act), _maxthreads(args._maxthreads), _nthreads(args._maxthreads),
683 _k_block(get_k_block_size(args)), _x_block(get_x_block_size(args)), _Mround(roundup(args._Msize, strategy::out_height())),
684 _os(os) { }
685
686 /* Constructor without OutputStage */
Georgios Pinitas48b3ef82019-10-14 19:03:09 +0100687 GemmInterleaved(const GemmArgs &args)
Georgios Pinitascfa2bba2019-06-27 17:00:52 +0100688 : _ci(args._ci), _Msize(args._Msize), _Nsize(args._Nsize), _Ksize(args._Ksize),
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000689 _Ksections(args._Ksections), _Ktotal(get_ktotal(args)),
690 _rounded_Ksize(roundup(_Ksize, strategy::k_unroll())),
691 _nbatches(args._nbatches), _nmulti(args._nmulti), _thread_columns(is_thread_columns(args)),
David Mansell318c9f42020-07-08 13:28:45 +0100692 _act(args._act), _maxthreads(args._maxthreads), _nthreads(args._maxthreads),
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000693 _k_block(get_k_block_size(args)), _x_block(get_x_block_size(args)), _Mround(roundup(args._Msize, strategy::out_height())),
694 _os() { }
Pablo Telloeb82fd22018-02-23 13:43:50 +0000695
696 // Interface implementation - Compulsory functions
697
Michalis Spyroue7e96e02018-04-13 13:44:10 +0100698 // Window size: Only the last thread should do a ragged block, so dole
699 // out work in units of out_height. Factor batches into the window, but
700 // not multi for now (as this would cause problems with the buffer
701 // manager).
Joseph Dobson6f8b17d2020-02-11 19:32:11 +0000702 ndrange_t get_window_size() const override {
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000703 unsigned int row_blocks = (_Mround / strategy::out_height()) * _nbatches;
704
705 if (_thread_columns) {
706 return { row_blocks, iceildiv(_Nsize, strategy::out_width()) };
707 } else {
708 // _Mround is a multiple of out_height by definition.
709 return { row_blocks };
710 }
Pablo Telloeb82fd22018-02-23 13:43:50 +0000711 }
712
713 // set_nthreads: pass on to buffer manager to avoid it waiting for non-existant threads.
Anthony Barbier5f707732018-07-03 16:22:02 +0100714 void set_nthreads(int nthreads) override {
715 _nthreads = std::min(nthreads, _maxthreads);
Pablo Telloeb82fd22018-02-23 13:43:50 +0000716 }
717
718 // Execute
Georgios Pinitas5aa1a0b2020-07-02 20:02:20 +0100719 void execute(const ndcoord_t &work_range, const ndcoord_t &, int threadid) override {
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100720#ifdef CYCLE_PROFILING
721 profiler prof;
722#endif
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100723
724 /* Make sure we've been set up correctly. */
Pablo Marquez Tello93581a52022-07-21 13:55:27 +0100725 assert(FixedFormat || _B_transposed);
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100726 assert(_working_space);
727 int8_t *working_space_bytes = reinterpret_cast<int8_t *>(_working_space);
728
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000729 /* Align if needed */
730 intptr_t working_space_v = reinterpret_cast<intptr_t>(_working_space);
731 if (working_space_v & 0x3f) {
732 intptr_t alignment_offset = 0x40 - (working_space_v & 0x3f);
733 working_space_bytes += alignment_offset;
734 }
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100735
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000736 strategy strat(_ci);
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100737
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000738 const auto start = work_range.get_position(0);
739 const auto end = work_range.get_position_end(0);
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100740
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000741 /* Translate 'start' and 'end' into a position within the batches and rows. */
742 const unsigned int window_per_batch = _Mround / strategy::out_height();
743 unsigned int batch_0 = start / window_per_batch;
744 unsigned int batch_end = end / window_per_batch;
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100745
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000746 // In ThreadColumns mode, process work one horizontal strip at a time.
747 // Transpose the block of needed rows at the start, then do all the work on that block.
748 if (_thread_columns) {
749 const auto start_x = work_range.get_position(1) * strategy::out_width();
750 const auto end_x = std::min(work_range.get_position_end(1) * strategy::out_width(), _Nsize);
751
752 Tri * const c_panel = reinterpret_cast<Tri *>(working_space_bytes + (threadid * get_c_working_size()));
753 Toi * const a_panel = reinterpret_cast<Toi *>(working_space_bytes + (_maxthreads * get_c_working_size()) +
754 (threadid * sizeof(Toi) * get_total_k_depth() * strategy::out_height()));
755
756 for (unsigned int multi=0; multi<_nmulti; multi++) {
757 for (unsigned int k0=0; k0<_Ktotal; k0+=_k_block) {
758 unsigned int kmax=std::min(k0+_k_block, _Ktotal);
759
760 unsigned int rounded_width = roundup(_Nsize, strategy::out_width());
761
762 const bool first_pass = (k0==0);
763 const bool last_pass = (kmax==_Ktotal);
764
765 // Figure out how many "K" the kernel will actually process.
766 unsigned int kern_k = roundup(kmax - k0, strategy::k_unroll());
767
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000768 const Toi *b_ptr = FixedFormat ?
769 reinterpret_cast<const Toi *>(this->_Bptr) + (multi * this->_B_multi_stride) +
770 ((start_x / get_stripe_width<strategy, FixedFormat>::get()) * this->_ldb) +
771 (k0 * get_stripe_width<strategy, FixedFormat>::get()) :
772 _B_transposed + (rounded_width * _Ktotal * multi) + (k0 * rounded_width) + (start_x * kern_k);
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000773
774 unsigned int batch = batch_0;
775 unsigned int start_row = (start - (batch_0 * window_per_batch)) * strategy::out_height();
776
777 for (unsigned int p=start; p<end; p++) {
778 unsigned int end_row = std::min(start_row + strategy::out_height(), _Msize);
779
780 // Set up transposed 'A' block
781 {
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100782#ifdef CYCLE_PROFILING
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000783 auto p=prof.ScopedProfiler(PROFILE_PREPA, strategy::out_height() * (kmax-k0) * sizeof(Toi));
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100784#endif
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000785 // See comment above on transform_type<> class: this extracts either 'transforms' or
786 // 'transforms_quantized' as appropriate.
787 typename transform_type<strategy, MergeStep && std::is_same<OutputStage, Requantize32>::value>::type transforms;
788
789 if (_indirect_buf != nullptr) {
790 transforms.PrepareA_indirect(a_panel,
791 _indirect_buf + (multi * _nbatches * _Ksections) + (batch * _Ksections), _Ksize,
792 _rounded_Ksize, start_row, end_row, k0, kmax, row_sum_multiplier());
793 } else if (_convolver) {
794 transforms.PrepareA_convolution(a_panel,
795 this->_Aptr + (batch * this->_A_batch_stride) + (multi * this->_A_multi_stride),
796 this->_lda, *_convolver, _rounded_Ksize, start_row, end_row, k0, kmax, row_sum_multiplier());
797 } else {
798 transforms.PrepareA(a_panel,
799 this->_Aptr + (batch * this->_A_batch_stride) + (multi * this->_A_multi_stride),
800 this->_lda, start_row, end_row, k0, std::min(kmax, _Ksize), row_sum_multiplier());
801 }
802 }
803
804 // Perform the kernel and merge step, either separately or together as required.
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000805 kernel_and_merge<MergeStep, FixedFormat, OutputStage>::run(
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000806 #ifdef CYCLE_PROFILING
807 prof,
808 #endif
809 // Strategy and panel pointers
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000810 strat, a_panel, b_ptr, this->_ldb, c_panel,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000811 // Result buffer pointers
812 this->_Cptr + (batch * this->_C_batch_stride) + (multi * this->_C_multi_stride), this->_ldc,
813 // K size, and M/N ranges
814 kern_k, start_row, end_row, start_x, end_x,
815 // Only do bias on the first pass
816 ((first_pass && this->_bias) ? this->_bias + (multi * this->_bias_multi_stride) : nullptr),
817 // Only do activation on the last pass, and accumulation on any non-first pass.
818 (last_pass ? _act : Activation()), !first_pass,
819 // Pass in quantization parameters for requantizing kernels (others will ignore)
820 _os, col_bias + (multi * _Nsize),
821 // Accumulation buffer (not yet implemented on this path)
822 static_cast<Tab *>(nullptr));
823
824 /* Increment to the next block */
825 start_row += strategy::out_height();
826 if (start_row >= _Msize) {
827 start_row = 0;
828 batch++;
829 }
830 }
831 }
832 }
833 } else {
834 blockwalker current(*this);
835
836 /* Compute the M values to operate on */
837 unsigned int m_0 = (start - (batch_0 * window_per_batch)) * strategy::out_height();
838 unsigned int m_max = (end - (batch_end * window_per_batch)) * strategy::out_height();
839
840 // Private buffers. Treat working_space as an array of C buffers
841 // (one per thread) first, followed by the (window-divided) A
842 // buffer.
843 // Set a_panel to the base of the A buffers - compute offsets into it based on M/batches later.
844 Toi * const a_panel = reinterpret_cast<Toi *>(working_space_bytes + (_maxthreads * get_c_working_size()));
845 Tri * const c_panel = reinterpret_cast<Tri *>(working_space_bytes + (threadid * get_c_working_size()));
846
847 const Toi *b_panel;
848 b_panel = _B_transposed;
849
850 // newkblock() is always true on the first iteration, so these will be set properly on the first loop.
851
852 // kern_k tracks the accumulation depth for the CURRENT K block a_panel_stride similarly tracks the total
853 // stride of the A panel (i.e. with 4 added for cases with embedded row sums)
854
855 // These are distinct from k_block and get_total_k_depth() which are based on the target K block size, and
856 // used for addressing inside a_panel.
857
858 // In cases where K blocking is in use and the blocks are not all the same size, the (smaller) final block
859 // won't use all the memory allocated.
860 unsigned int kern_k = 0;
861 unsigned int a_panel_stride = 0;
862
863 for (;!current.done();current.advance()) {
864 if (current.newkblock()) {
865#ifdef CYCLE_PROFILING
866 auto p=prof.ScopedProfiler(PROFILE_PREPA, (end - start) * strategy::out_height() * (current.kmax()-current.k0()) * sizeof(Toi));
867#endif
868 // See comment above on transform_type<> class: this extracts either 'transforms' or
869 // 'transforms_quantized' as appropriate.
870 typename transform_type<strategy, MergeStep && std::is_same<OutputStage, Requantize32>::value>::type transforms;
871
872 for (unsigned int batch = batch_0; batch <= batch_end; batch++) {
873 unsigned int first_m = (batch == batch_0) ? m_0 : 0;
874 unsigned int last_m = (batch == batch_end) ? m_max : _Msize;
875
876 if (first_m >= last_m)
877 continue;
878
879 if (_indirect_buf != nullptr) {
880 transforms.PrepareA_indirect(a_panel + ((batch * _Mround + first_m) * get_total_k_depth()),
881 _indirect_buf + (current.multi() * _nbatches * _Ksections) + (batch * _Ksections), _Ksize,
882 _rounded_Ksize, first_m, last_m, current.k0(), current.kmax(), row_sum_multiplier());
883 } else if (_convolver) {
884 transforms.PrepareA_convolution(a_panel + ((batch * _Mround + first_m) * get_total_k_depth()),
885 this->_Aptr + (batch * this->_A_batch_stride) + (current.multi() * this->_A_multi_stride),
886 this->_lda, *_convolver, _rounded_Ksize, first_m, last_m, current.k0(), current.kmax(), row_sum_multiplier());
887 } else {
888 transforms.PrepareA(a_panel + ((batch * _Mround + first_m) * get_total_k_depth()),
889 this->_Aptr + (batch * this->_A_batch_stride) + (current.multi() * this->_A_multi_stride),
890 this->_lda, first_m, last_m, current.k0(), std::min(_Ksize, current.kmax()), row_sum_multiplier());
891 }
892 }
893
894 // Figure out how many "K" the kernel will actually process.
895 kern_k = roundup(current.kmax() - current.k0(), strategy::k_unroll());
896
897 // Requantizing GEMMs have the row sums built in to the
898 // transposed data, so the stride between rows is 4 bytes
899 // larger than the (rounded) K value.
900
901 if(std::is_same<OutputStage, Requantize32>::value) {
902 a_panel_stride = kern_k + (sizeof(int32_t) / sizeof(Toi));
903 } else {
904 a_panel_stride = kern_k;
905 }
906 }
907
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000908 // For FixedFormat cases, figure out the B pointer. The loop below moves through batches and vertically through the output so this will be the same throughout.
909 if (FixedFormat) {
910 b_panel = reinterpret_cast<const Toi *>(this->_Bptr) + (current.multi() * this->_B_multi_stride) +
911 ((current.x0() / get_stripe_width<strategy, FixedFormat>::get()) * this->_ldb) +
912 (current.k0() * get_stripe_width<strategy, FixedFormat>::get());
913 }
914
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000915 /* Do the actual work. */
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100916 for (unsigned int batch = batch_0; batch <= batch_end; batch++) {
917 unsigned int first_m = (batch == batch_0) ? m_0 : 0;
918 unsigned int last_m = (batch == batch_end) ? m_max : _Msize;
919
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000920 const Toi *a_ptr = a_panel + (batch * _Mround + first_m) * get_total_k_depth();
921
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100922 if (first_m >= last_m)
923 continue;
924
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000925 // For the merge case we need to do this out_height() rows
926 // at a time, as that is the size of our intermediate
927 // buffer. If we are not doing that, we can do all the
928 // relevant rows in one go.
929 unsigned int m_step = MergeStep ? strategy::out_height() : (last_m - first_m);
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100930
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000931 // But in the case where we have an accumulation buffer, we can't do that after all, unless
932 // there is no N blocking.
933 if (_accumulation_buffer && ((current.x0() != 0) || (current.xmax() < _Nsize))) {
934 m_step = strategy::out_height();
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100935 }
936
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000937 for (unsigned int y=first_m; y<last_m; y+=m_step) {
938 unsigned int ymax = std::min(_Msize, y + m_step);
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100939
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000940 const bool first_pass = (current.k0() == 0);
941 const bool last_pass = (current.kmax() == _Ktotal);
942
943 // Pointer to appropriate part of result array.
944 Tr *result_ptr = this->_Cptr + (batch * this->_C_batch_stride) + (current.multi() * this->_C_multi_stride);
945
946 // If we are using an accumulation buffer, we don't pass the result buffer to ask the kernel
947 // to write things into the accumulation buffer instead, except on the last pass.
948 if (_accumulation_buffer && !last_pass) {
949 result_ptr = nullptr;
950 }
951
952 // Perform the kernel and merge step, either separately or together as required.
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000953 kernel_and_merge<MergeStep, FixedFormat, OutputStage>::run(
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000954 #ifdef CYCLE_PROFILING
955 prof,
956 #endif
957 // Strategy and panel pointers
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000958 strat, a_ptr, b_panel, this->_ldb, c_panel,
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000959 // Result buffer pointers
960 result_ptr, this->_ldc,
961 // K size, and M/N ranges
962 kern_k, y, ymax, current.x0(), current.xmax(),
963 // Only do bias on the first pass
964 ((first_pass && this->_bias) ? this->_bias + (current.multi() * this->_bias_multi_stride) : nullptr),
965 // Only do activation on the last pass, and accumulation on any non-first pass.
966 (last_pass ? _act : Activation()), !first_pass,
967 // Pass in quantization parameters for requantizing kernels (others will ignore)
968 _os, col_bias + (current.multi() * _Nsize),
969 // Accumulation buffer
970 get_accumulation_buffer(y, current.x0(), batch, current.multi()) );
971
972 a_ptr += (strategy::out_height() * a_panel_stride);
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100973 }
974 }
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +0100975
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +0000976 if (FixedFormat == false) {
977 b_panel += (roundup(current.xmax() - current.x0(), strategy::out_width()) * kern_k);
978 }
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000979 }
Pablo Telloeb82fd22018-02-23 13:43:50 +0000980 }
981 }
982
983 // Interface implementation - working space
Anthony Barbier5f707732018-07-03 16:22:02 +0100984 size_t get_working_size() const override {
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000985 // In all cases, we need one A buffer plus a C buffer per thread, plus an accumulation buffer.
986 size_t size = get_a_working_size() + (get_c_working_size() * _maxthreads) + get_accumulation_buffer_size();
Pablo Telloeb82fd22018-02-23 13:43:50 +0000987
Georgios Pinitasc0b6f762020-11-02 01:37:17 +0000988 size += 128; // Add on two cache lines extra for alignment.
Pablo Telloeb82fd22018-02-23 13:43:50 +0000989
990 return size;
991 }
992
Anthony Barbier5f707732018-07-03 16:22:02 +0100993 void set_working_space(void *working_space) override {
Pablo Telloeb82fd22018-02-23 13:43:50 +0000994 // Make sure everything ends up cache line aligned
995 int8_t *working_space_bytes = reinterpret_cast<int8_t *>(working_space);
Anthony Barbier5f707732018-07-03 16:22:02 +0100996 intptr_t working_space_int = reinterpret_cast<intptr_t>(working_space);
Pablo Telloeb82fd22018-02-23 13:43:50 +0000997
Anthony Barbier5f707732018-07-03 16:22:02 +0100998 size_t diff=0;
Pablo Telloeb82fd22018-02-23 13:43:50 +0000999
Anthony Barbier5f707732018-07-03 16:22:02 +01001000 if (working_space_int & 0x3F) {
Pablo Telloeb82fd22018-02-23 13:43:50 +00001001 diff = 0x40 - (working_space_int & 0x3F);
1002 }
1003
1004 working_space_bytes += diff;
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001005 working_space_int += diff;
Pablo Telloeb82fd22018-02-23 13:43:50 +00001006
Georgios Pinitas0cc50ed2020-07-06 19:10:38 +01001007 // Pretransposed case: just set internal pointer to parameter value.
1008 _working_space = reinterpret_cast<void *>(working_space_bytes);
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001009
1010 // Set up accumulation buffer
1011 if (get_accumulation_buffer_size() > 0) {
1012 intptr_t acc_buff_int = working_space_int + get_a_working_size() + (get_c_working_size() * _maxthreads);
1013 // Make sure the accumulation buffer is aligned (needed if the other blocks are not a multiple of cache line length)
1014 if (acc_buff_int & 0x3F) {
1015 acc_buff_int += (0x40 - (acc_buff_int & 0x3F));
1016 }
1017 _accumulation_buffer = reinterpret_cast<Tab *>(acc_buff_int);
1018 } else {
1019 _accumulation_buffer = nullptr;
1020 }
Pablo Telloeb82fd22018-02-23 13:43:50 +00001021 }
1022
1023 // Interface implementation - pretransposed
Anthony Barbier5f707732018-07-03 16:22:02 +01001024 bool B_is_pretransposed() const override {
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +00001025 return (FixedFormat == false);
Pablo Telloeb82fd22018-02-23 13:43:50 +00001026 }
1027
Anthony Barbier5f707732018-07-03 16:22:02 +01001028 bool B_pretranspose_required() const override {
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +00001029 return (FixedFormat == false) && (_B_transposed==nullptr);
Pablo Telloeb82fd22018-02-23 13:43:50 +00001030 }
1031
Anthony Barbier5f707732018-07-03 16:22:02 +01001032 size_t get_B_pretransposed_array_size() const override {
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +00001033 if (FixedFormat) {
1034 return 0;
1035 }
1036
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001037 unsigned int x_size = roundup(_Nsize, strategy::out_width());
Pablo Telloeb82fd22018-02-23 13:43:50 +00001038
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001039 return (x_size * _Ktotal * _nmulti * sizeof(Toi)) + get_col_sum_size();
Pablo Telloeb82fd22018-02-23 13:43:50 +00001040 }
1041
Giorgio Arena63e0beb2021-09-24 14:04:27 +01001042 void requantize_bias(void *in_buffer, const To *B, const int ldb, const int B_multi_stride) override {
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001043 if (std::is_same<OutputStage, Requantize32>::value) {
1044 col_bias = reinterpret_cast<int32_t *>(in_buffer);
1045
1046 Requantize32 *qp_ptr = reinterpret_cast<Requantize32 *>(&_os);
1047
1048 for (unsigned int i=0; i<_nmulti; i++) {
1049 // The input is assumed not to have any padding between sections, so straightforward Ksize * Ksections computation gets the total size.
1050 compute_col_sums(*qp_ptr, _Nsize, _Ksize * _Ksections, B + (i * B_multi_stride), ldb, col_bias + (i * _Nsize), _Ksize * _Ksections, i, 0);
1051 }
1052 }
Giorgio Arena63e0beb2021-09-24 14:04:27 +01001053 }
1054
1055 void pretranspose_B_array(void *in_buffer, const To *B, const int ldb, const int B_multi_stride) override {
1056 requantize_bias(in_buffer, B, ldb, B_multi_stride);
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001057
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +00001058 // Put the transposed data after the column sums - in non-quantized cases get_col_sum_size() == 0
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001059 uintptr_t buffer_int = reinterpret_cast<uintptr_t>(in_buffer);
1060 Toi *buffer = reinterpret_cast<Toi *>(buffer_int + get_col_sum_size());
Anthony Barbier5f707732018-07-03 16:22:02 +01001061 _B_transposed = buffer;
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001062
1063 blockwalker current(*this);
David Manselld93991e2018-07-06 14:52:52 +01001064 strategy strat(_ci);
Pablo Telloeb82fd22018-02-23 13:43:50 +00001065
Anthony Barbier5f707732018-07-03 16:22:02 +01001066 do {
Pablo Telloeb82fd22018-02-23 13:43:50 +00001067 /* Figure out the size of each block. */
Georgios Pinitas1d480652019-01-23 11:24:50 +00001068 unsigned int k_size = (current.kmax() - current.k0());
Pablo Telloeb82fd22018-02-23 13:43:50 +00001069
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001070 if (_Ksections > 1) {
1071 // We need to insert padding at the end of each K section.
1072 // The computation needed is a little delicate - the coordinates from the block walker are expressed in
1073 // terms of the full, padded, _Ktotal.
1074 // But we need to transform each section with reference to the original, unpadded, input, letting the
1075 // transform pad each section as needed.
Pablo Telloeb82fd22018-02-23 13:43:50 +00001076
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001077 // This is needed for computations below.
1078 const unsigned int rounded_section_size = roundup(_Ksize, strategy::k_unroll());
Pablo Telloeb82fd22018-02-23 13:43:50 +00001079
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001080 // The expected output format is also an entire <out_width> columns interleaved, then the next set of
1081 // columns, and so on. This means, as we are breaking it up vertically, we have to do it one column at
1082 // a time.
1083 for (unsigned int x0=current.x0(); x0 < current.xmax(); x0 += strategy::out_width() ) {
1084 unsigned int xmax = std::min(x0 + strategy::out_width(), current.xmax());
Pablo Telloeb82fd22018-02-23 13:43:50 +00001085
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001086 // Track where we are and how much work is left.
1087 unsigned int kpos = current.k0();
1088 unsigned int kleft = k_size;
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001089
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001090 while (kleft) {
1091 // Which section are we in? Based on the rounded-up section size.
1092 unsigned int k_section_base = kpos / rounded_section_size;
1093 // How far into the section are we?
1094 unsigned int k_offset = kpos - (k_section_base * rounded_section_size);
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001095
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001096 // We will either copy the rest of this section, or to the end of the requested length.
1097 unsigned int k_length = std::min(_Ksize - k_offset, kleft);
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001098
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001099 strat.transforms.PrepareB(buffer, B + (current.multi() * B_multi_stride), ldb,
1100 x0, xmax,
1101 (k_section_base * _Ksize) + k_offset, // K starting point - compute row to read based on our section and the true section length.
1102 (k_section_base * _Ksize) + k_offset + k_length); // K end point - starting point plus length computed above.
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001103
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001104 // We need to modify our position based on the ROUNDED version of what we just did.
1105 unsigned int padded_length = roundup(k_length, strategy::k_unroll());
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001106
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001107 buffer += strategy::out_width() * padded_length;
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001108
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001109 kpos += padded_length;
1110 kleft -= padded_length;
1111 }
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001112 }
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001113 } else {
1114 // In the single K section case, can process the whole lot in one go.
1115 // Caution: 'blockwalker::kmax()' rounds up, so clamp to valid _Ksize.
1116 strat.transforms.PrepareB(buffer, B + (current.multi() * B_multi_stride), ldb,
1117 current.x0(), current.xmax(), current.k0(), std::min(current.kmax(), _Ksize));
1118 buffer += roundup(current.xmax() - current.x0(), strategy::out_width()) * roundup(current.kmax() - current.k0(), strategy::k_unroll());
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001119 }
Anthony Barbier5f707732018-07-03 16:22:02 +01001120 } while (current.advance());
Pablo Telloeb82fd22018-02-23 13:43:50 +00001121 }
1122
Anthony Barbier5f707732018-07-03 16:22:02 +01001123 void set_pretransposed_B_data(void *in_buffer) override {
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +00001124 // Put the transposed data after the column sums - in non-quantized cases get_col_sum_size() == 0
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001125 uintptr_t buffer_int = reinterpret_cast<uintptr_t>(in_buffer);
1126 _B_transposed = reinterpret_cast<Toi *>(buffer_int + get_col_sum_size());
1127 col_bias = reinterpret_cast<int32_t *>(in_buffer);
1128 }
1129
1130 void set_quantized_bias(const int32_t *bias, size_t bias_multi_stride) override {
1131 if (std::is_same<OutputStage, Requantize32>::value) {
1132 Requantize32 *qp = reinterpret_cast<Requantize32 *>(&_os);
1133
1134 qp->bias = bias;
1135 qp->bias_multi_stride = bias_multi_stride;
1136 }
1137 }
1138
1139 void set_indirect_parameters(size_t string_len, const To * const * const *ptr) override {
1140 assert(string_len == _Ksize);
1141 _indirect_buf = ptr;
1142 }
1143
1144 void set_convolution_parameters(ConvolutionParameters parms) override {
1145 assert(parms.input_channels == _Ksize);
1146 _convolver = std::unique_ptr<convolver<To>>(new convolver<To>(parms));
Michalis Spyroue7e96e02018-04-13 13:44:10 +01001147 }
David Mansell318c9f42020-07-08 13:28:45 +01001148
1149 // Estimate cycles for given problem given provided parameters
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001150 template<typename perf_type>
1151 static uint64_t estimate_cycles(const GemmArgs &args) {
David Mansell318c9f42020-07-08 13:28:45 +01001152 unsigned int k_blocks = iceildiv(args._Ksize, get_k_block_size(args));
1153
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001154 const PerformanceParameters &params = strategy::template get_performance_parameters<perf_type>(args._ci);
1155
Georgios Pinitas6f45cf72021-02-23 23:41:40 +00001156 uint64_t total_macs = static_cast<uint64_t>(args._nbatches) * args._nmulti * roundup(args._Msize, strategy::out_height()) * roundup(args._Nsize, strategy::out_width()) * get_ktotal(args);
1157 uint64_t prepare_bytes = static_cast<uint64_t>(args._nbatches) * args._nmulti * roundup(args._Msize, strategy::out_height()) * get_ktotal(args) * sizeof(Toi);
ramelg011f864492022-07-07 15:12:20 +01001158 uint64_t merge_bytes = static_cast<uint64_t>(args._nbatches) * args._nmulti * k_blocks * args._Msize * roundup(args._Nsize, strategy::out_width()) * sizeof(Tr);
David Mansell318c9f42020-07-08 13:28:45 +01001159
1160 float mac_cycles = static_cast<float>(total_macs) / params.kernel_macs_cycle;
1161 float prepare_cycles = static_cast<float>(prepare_bytes) / params.prepare_bytes_cycle;
1162 float merge_cycles = static_cast<float>(merge_bytes) / params.merge_bytes_cycle;
1163
1164 float total_cycles = mac_cycles + prepare_cycles + merge_cycles;
1165
1166 // We can't thread over multis or width, which makes this a poor
1167 // choice in many threaded cases. Penalize that here.
1168 float parallelism_available = static_cast<float>(iceildiv(args._Msize, strategy::out_height()) * args._nbatches) * 0.9f;
1169
1170 if (parallelism_available < args._maxthreads) {
1171 total_cycles *= (static_cast<float>(args._maxthreads) / parallelism_available);
1172 }
1173
1174 return static_cast<uint64_t>(total_cycles);
1175 }
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001176
1177 GemmConfig get_config() override {
1178 GemmConfig c;
1179
1180 c.method = GemmMethod::GEMM_INTERLEAVED;
1181 c.inner_block_size = _k_block;
1182 c.outer_block_size = _x_block;
1183 c.filter = get_type_name<strategy>();
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +00001184 c.weight_format = get_weight_format(get_kernel_weight_format<strategy, FixedFormat, To>::get(), sizeof(To));
Georgios Pinitas4ee8b152021-07-16 16:16:43 +01001185
1186 return c;
1187 }
Pablo Telloeb82fd22018-02-23 13:43:50 +00001188};
1189
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001190// Aliases for the variations
1191template<typename strategy, typename To, typename Tr, typename OutputStage=Nothing>
1192using GemmInterleavedNoMerge = GemmInterleaved<strategy, To, Tr, OutputStage, false>;
1193
Francesco.Petrogalli@arm.com5fcf22d2022-04-05 10:31:08 +00001194template<typename strategy, typename To, typename Tr, typename OutputStage=Nothing>
1195using GemmInterleavedFixedFormat = GemmInterleaved<strategy, To, Tr, OutputStage, true, true>;
1196
Georgios Pinitasc0b6f762020-11-02 01:37:17 +00001197template<typename strategy, typename To, typename Tr>
1198using GemmInterleavedPretransposedNoMergeQuantizedInline = GemmInterleaved<strategy, To, Tr, Requantize32, false>;
1199
1200template<typename strategy, typename To, typename Tr>
1201using GemmInterleavedQuantized = GemmInterleaved<strategy, To, Tr, Requantize32>;
1202
Pablo Telloeb82fd22018-02-23 13:43:50 +00001203} // namespace arm_gemm