xref: /aosp_15_r20/external/ComputeLibrary/src/core/NEON/kernels/arm_gemm/kernels/a64_sgemm_8x12/generic.cpp (revision c217d954acce2dbc11938adb493fc0abd69584f3)
1 /*
2  * Copyright (c) 2017-2018 Arm Limited.
3  *
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 #ifdef __aarch64__
25 
26 #include <arm_neon.h>
27 
28 #include "../../asmlib.hpp"
29 
30 // Kernel implementation.
31 //
32 // Assume that "Apanel" points to a chunk of A blocks (each size 8xK) in read-order.
33 // Assume that "Bpanel" points to a chunk of B blocks (each size 12xK) in read-order.
34 // Assume that "Cpanel" points to a chunk of C output blocks (each size
35 // 12x8), the chunks being arranged in a row major fashion.
36 //
37 // Note that the intent of this is that either ablocks or bblocks will be 1
38 // - this construction allows the output loop to proceed in either order.
39 
40 namespace arm_gemm {
41 
a64_sgemm_asimd_8x12(const float * Apanel,const float * Bpanel,float * Cpanel,int ablocks,int bblocks,int K)42 void a64_sgemm_asimd_8x12(const float *Apanel, const float *Bpanel, float *Cpanel, int ablocks, int bblocks, int K) {
43     const float *a_ptr = Apanel;
44     float *c_ptr = Cpanel;
45 
46     for (int yb=0; yb<ablocks; yb++) {
47         const float *a_ptr0 = a_ptr;
48         const float *b_ptr = Bpanel;
49 
50         for (int xb=0; xb<bblocks; xb++) {
51             a_ptr = a_ptr0;
52             // Fix up for odd lengths - set a flag if K is odd, but make
53             // sure we round up the iteration count.
54             int oddk = (K & 1);
55             int k = ((K+1)/2) - 1;
56 
57             register float32x4_t a0  asm("v0");
58             register float32x4_t a1  asm("v1");
59             register float32x4_t b0  asm("v2");
60             register float32x4_t b1  asm("v3");
61             register float32x4_t b2  asm("v4");
62             register float32x4_t a0a asm("v5");
63             register float32x4_t a1a asm("v6");
64 
65             __asm __volatile (
66                 // Initialize result registers, load initial operands, prime prefetches.
67                 "movi	v8.4s, #0x0\n"
68                 "ldr	%q[a0], [%[a_ptr]]\n"
69                 "movi	v9.4s, #0x0\n"
70                 "ldr	%q[b0], [%[b_ptr]]\n"
71                 "movi	v10.4s, #0x0\n"
72                 "ldr	%q[a1], [%[a_ptr], #16]\n"
73                 "movi	v11.4s, #0x0\n"
74                 "ldr	%q[b1], [%[b_ptr], #16]\n"
75                 "movi	v12.4s, #0x0\n"
76                 ASM_PREFETCH("[%[b_ptr], #64]")
77                 "movi	v13.4s, #0x0\n"
78                 ASM_PREFETCH("[%[a_ptr], #64]")
79                 "movi	v14.4s, #0x0\n"
80                 ASM_PREFETCH("[%[b_ptr], #128]")
81                 "movi	v15.4s, #0x0\n"
82                 ASM_PREFETCH("[%[a_ptr], #128]")
83                 "movi	v16.4s, #0x0\n"
84                 ASM_PREFETCH("[%[b_ptr], #192]")
85                 "movi	v17.4s, #0x0\n"
86                 ASM_PREFETCH("[%[b_ptr], #256]")
87                 "movi	v18.4s, #0x0\n"
88                 ASM_PREFETCH("[%[a_ptr], #192]")
89                 "movi	v19.4s, #0x0\n"
90                 ASM_PREFETCH("[%[b_ptr], #320]")
91                 "movi	v20.4s, #0x0\n"
92                 ASM_PREFETCH("[%[a_ptr], #256]")
93                 "movi	v21.4s, #0x0\n"
94                 ASM_PREFETCH("[%[b_ptr], #384]")
95                 "movi	v22.4s, #0x0\n"
96                 "movi	v23.4s, #0x0\n"
97                 "movi	v24.4s, #0x0\n"
98                 "movi	v25.4s, #0x0\n"
99                 "movi	v26.4s, #0x0\n"
100                 "movi	v27.4s, #0x0\n"
101                 "movi	v28.4s, #0x0\n"
102                 "movi	v29.4s, #0x0\n"
103                 "movi	v30.4s, #0x0\n"
104                 "movi	v31.4s, #0x0\n"
105 
106                 // Skip loop if we are doing zero iterations of it.
107                 "cbz	%w[k], 4f\n"
108 
109                 // Loop proper
110                 "1:\n"
111                 "fmla 	v8.4s , %[b0].4s, %[a0].s[0]\n"
112                 "fmla  	v9.4s , %[b0].4s, %[a0].s[1]\n"
113                 "ldr	%q[b2], [%[b_ptr], #32]\n"
114                 "fmla	v10.4s, %[b0].4s, %[a0].s[2]\n"
115                 "fmla	v11.4s, %[b0].4s, %[a0].s[3]\n"
116                 "ldr	%q[a0a], [%[a_ptr], #32]\n"
117                 "fmla 	v12.4s, %[b0].4s, %[a1].s[0]\n"
118                 "fmla	v13.4s, %[b0].4s, %[a1].s[1]\n"
119                 "ldr	%q[a1a], [%[a_ptr], #48]\n"
120                 "fmla	v14.4s, %[b0].4s, %[a1].s[2]\n"
121                 "fmla	v15.4s, %[b0].4s, %[a1].s[3]\n"
122                 "ldr	%q[b0], [%[b_ptr], #48]\n"
123 
124                 "fmla	v16.4s, %[b1].4s, %[a0].s[0]\n"
125                 "fmla	v17.4s, %[b1].4s, %[a0].s[1]\n"
126                 ASM_PREFETCH("[%[a_ptr], #320]")
127                 "fmla	v18.4s, %[b1].4s, %[a0].s[2]\n"
128                 "fmla	v19.4s, %[b1].4s, %[a0].s[3]\n"
129                 "fmla	v20.4s, %[b1].4s, %[a1].s[0]\n"
130                 "fmla	v21.4s, %[b1].4s, %[a1].s[1]\n"
131                 "fmla	v22.4s, %[b1].4s, %[a1].s[2]\n"
132                 "fmla	v23.4s, %[b1].4s, %[a1].s[3]\n"
133                 "ldr	%q[b1], [%[b_ptr], #64]\n"
134 
135                 "fmla	v24.4s, %[b2].4s, %[a0].s[0]\n"
136                 "fmla	v25.4s, %[b2].4s, %[a0].s[1]\n"
137                 ASM_PREFETCH("[%[b_ptr], #448]")
138                 "fmla	v26.4s, %[b2].4s, %[a0].s[2]\n"
139                 "fmla	v27.4s, %[b2].4s, %[a0].s[3]\n"
140                 "fmla	v28.4s, %[b2].4s, %[a1].s[0]\n"
141                 "fmla	v29.4s, %[b2].4s, %[a1].s[1]\n"
142                 "fmla	v30.4s, %[b2].4s, %[a1].s[2]\n"
143                 "fmla	v31.4s, %[b2].4s, %[a1].s[3]\n"
144                 "ldr	%q[b2], [%[b_ptr], #80]\n"
145 
146                 "fmla 	v8.4s , %[b0].4s, %[a0a].s[0]\n"
147                 "fmla	v9.4s , %[b0].4s, %[a0a].s[1]\n"
148                 "ldr	%q[a0], [%[a_ptr], #64]\n"
149                 "fmla	v10.4s, %[b0].4s, %[a0a].s[2]\n"
150                 "fmla	v11.4s, %[b0].4s, %[a0a].s[3]\n"
151                 "fmla 	v12.4s, %[b0].4s, %[a1a].s[0]\n"
152                 "ldr	%q[a1], [%[a_ptr], #80]\n"
153                 "fmla   v13.4s, %[b0].4s, %[a1a].s[1]\n"
154                 "fmla	v14.4s, %[b0].4s, %[a1a].s[2]\n"
155                 "fmla	v15.4s, %[b0].4s, %[a1a].s[3]\n"
156                 "ldr	%q[b0], [%[b_ptr], #96]\n"
157 
158                 "fmla	v16.4s, %[b1].4s, %[a0a].s[0]\n"
159                 "fmla	v17.4s, %[b1].4s, %[a0a].s[1]\n"
160                 ASM_PREFETCH("[%[b_ptr], #512]")
161                 "fmla	v18.4s, %[b1].4s, %[a0a].s[2]\n"
162                 "fmla	v19.4s, %[b1].4s, %[a0a].s[3]\n"
163                 "fmla	v20.4s, %[b1].4s, %[a1a].s[0]\n"
164                 "fmla	v21.4s, %[b1].4s, %[a1a].s[1]\n"
165                 "fmla	v22.4s, %[b1].4s, %[a1a].s[2]\n"
166                 "fmla	v23.4s, %[b1].4s, %[a1a].s[3]\n"
167                 "ldr	%q[b1], [%[b_ptr], #112]\n"
168 
169                 "fmla	v24.4s, %[b2].4s, %[a0a].s[0]\n"
170                 "fmla	v25.4s, %[b2].4s, %[a0a].s[1]\n"
171                 "add	%[a_ptr], %[a_ptr], #64\n"
172                 "fmla	v26.4s, %[b2].4s, %[a0a].s[2]\n"
173                 "fmla	v27.4s, %[b2].4s, %[a0a].s[3]\n"
174                 "add	%[b_ptr], %[b_ptr], #96\n"
175                 "fmla	v28.4s, %[b2].4s, %[a1a].s[0]\n"
176                 "fmla	v29.4s, %[b2].4s, %[a1a].s[1]\n"
177                 "subs	%w[k], %w[k], #1\n"
178                 "fmla	v30.4s, %[b2].4s, %[a1a].s[2]\n"
179                 "fmla	v31.4s, %[b2].4s, %[a1a].s[3]\n"
180                 "bne	1b\n"
181 
182                 // Target to use when K is 1 or 2 (i.e. zero iterations of main loop)
183                 "4:\n"
184 
185                 // Branch to alternative tail for odd K
186                 "cbnz	%w[oddk], 2f\n"
187 
188                 // Detached final iteration (even K)
189                 "fmla 	v8.4s , %[b0].4s, %[a0].s[0]\n"
190                 "fmla   v9.4s , %[b0].4s, %[a0].s[1]\n"
191                 "ldr	%q[b2], [%[b_ptr], #32]\n"
192                 "fmla	v10.4s, %[b0].4s, %[a0].s[2]\n"
193                 "fmla	v11.4s, %[b0].4s, %[a0].s[3]\n"
194                 "ldr	%q[a0a], [%[a_ptr], #32]\n"
195                 "fmla 	v12.4s, %[b0].4s, %[a1].s[0]\n"
196                 "fmla   v13.4s, %[b0].4s, %[a1].s[1]\n"
197                 "ldr	%q[a1a], [%[a_ptr], #48]\n"
198                 "fmla	v14.4s, %[b0].4s, %[a1].s[2]\n"
199                 "fmla	v15.4s, %[b0].4s, %[a1].s[3]\n"
200                 "ldr	%q[b0], [%[b_ptr], #48]\n"
201 
202                 "fmla	v16.4s, %[b1].4s, %[a0].s[0]\n"
203                 "fmla	v17.4s, %[b1].4s, %[a0].s[1]\n"
204                 "fmla	v18.4s, %[b1].4s, %[a0].s[2]\n"
205                 "fmla	v19.4s, %[b1].4s, %[a0].s[3]\n"
206                 "fmla	v20.4s, %[b1].4s, %[a1].s[0]\n"
207                 "fmla	v21.4s, %[b1].4s, %[a1].s[1]\n"
208                 "fmla	v22.4s, %[b1].4s, %[a1].s[2]\n"
209                 "fmla	v23.4s, %[b1].4s, %[a1].s[3]\n"
210                 "ldr	%q[b1], [%[b_ptr], #64]\n"
211 
212                 "fmla	v24.4s, %[b2].4s, %[a0].s[0]\n"
213                 "fmla	v25.4s, %[b2].4s, %[a0].s[1]\n"
214                 "add	%[a_ptr], %[a_ptr], #64\n"
215                 "fmla	v26.4s, %[b2].4s, %[a0].s[2]\n"
216                 "fmla	v27.4s, %[b2].4s, %[a0].s[3]\n"
217                 "fmla	v28.4s, %[b2].4s, %[a1].s[0]\n"
218                 "fmla	v29.4s, %[b2].4s, %[a1].s[1]\n"
219                 "fmla	v30.4s, %[b2].4s, %[a1].s[2]\n"
220                 "fmla	v31.4s, %[b2].4s, %[a1].s[3]\n"
221                 "ldr	%q[b2], [%[b_ptr], #80]\n"
222 
223                 "fmla 	v8.4s , %[b0].4s, %[a0a].s[0]\n"
224                 "fmla	v16.4s, %[b1].4s, %[a0a].s[0]\n"
225                 "add	%[b_ptr], %[b_ptr], #96\n"
226                 "fmla   v9.4s , %[b0].4s, %[a0a].s[1]\n"
227                 "str	q8, [%[c_ptr], #0]\n"
228                 "fmla	v17.4s, %[b1].4s, %[a0a].s[1]\n"
229                 "str	q16, [%[c_ptr], #16]\n"
230                 "fmla	v24.4s, %[b2].4s, %[a0a].s[0]\n"
231                 "str	q24, [%[c_ptr], #32]\n"
232 
233                 "fmla	v25.4s, %[b2].4s, %[a0a].s[1]\n"
234                 "str	q9, [%[c_ptr], #48]\n"
235                 "fmla	v10.4s, %[b0].4s, %[a0a].s[2]\n"
236                 "str	q17, [%[c_ptr], #64]\n"
237                 "fmla	v18.4s, %[b1].4s, %[a0a].s[2]\n"
238                 "str	q25, [%[c_ptr], #80]\n"
239                 "fmla	v26.4s, %[b2].4s, %[a0a].s[2]\n"
240                 "str	q10, [%[c_ptr], #96]\n"
241 
242                 "fmla	v11.4s, %[b0].4s, %[a0a].s[3]\n"
243                 "str	q18, [%[c_ptr], #112]\n"
244                 "fmla	v19.4s, %[b1].4s, %[a0a].s[3]\n"
245                 "str	q26, [%[c_ptr], #128]\n"
246                 "fmla	v27.4s, %[b2].4s, %[a0a].s[3]\n"
247                 "str	q11, [%[c_ptr], #144]\n"
248 
249                 "fmla 	v12.4s, %[b0].4s, %[a1a].s[0]\n"
250                 "str	q19, [%[c_ptr], #160]\n"
251                 "fmla	v20.4s, %[b1].4s, %[a1a].s[0]\n"
252                 "str	q27, [%[c_ptr], #176]\n"
253                 "fmla	v28.4s, %[b2].4s, %[a1a].s[0]\n"
254                 "str	q12, [%[c_ptr], #192]\n"
255 
256                 "fmla   v13.4s, %[b0].4s, %[a1a].s[1]\n"
257                 "str	q20, [%[c_ptr], #208]\n"
258                 "fmla	v21.4s, %[b1].4s, %[a1a].s[1]\n"
259                 "str	q28, [%[c_ptr], #224]\n"
260                 "fmla	v29.4s, %[b2].4s, %[a1a].s[1]\n"
261                 "str	q13, [%[c_ptr], #240]\n"
262 
263                 "fmla	v14.4s, %[b0].4s, %[a1a].s[2]\n"
264                 "str	q21, [%[c_ptr], #256]\n"
265                 "fmla	v22.4s, %[b1].4s, %[a1a].s[2]\n"
266                 "str	q29, [%[c_ptr], #272]\n"
267                 "fmla	v30.4s, %[b2].4s, %[a1a].s[2]\n"
268                 "str	q14, [%[c_ptr], #288]\n"
269 
270                 "fmla	v15.4s, %[b0].4s, %[a1a].s[3]\n"
271                 "str	q22, [%[c_ptr], #304]\n"
272                 "fmla	v23.4s, %[b1].4s, %[a1a].s[3]\n"
273                 "str	q30, [%[c_ptr], #320]\n"
274                 "fmla	v31.4s, %[b2].4s, %[a1a].s[3]\n"
275                 "str	q15, [%[c_ptr], #336]\n"
276 
277                 "b	3f\n"
278 
279                 // Detached final iteration (odd K)
280                 "2:\n"
281                 "fmla 	v8.4s , %[b0].4s, %[a0].s[0]\n"
282                 "ldr	%q[b2], [%[b_ptr], #32]\n"
283                 "fmla	v16.4s, %[b1].4s, %[a0].s[0]\n"
284                 "fmla   v9.4s , %[b0].4s, %[a0].s[1]\n"
285                 "str	q8, [%[c_ptr], #0]\n"
286                 "fmla	v17.4s, %[b1].4s, %[a0].s[1]\n"
287                 "str	q16, [%[c_ptr], #16]\n"
288                 "fmla	v24.4s, %[b2].4s, %[a0].s[0]\n"
289                 "add	%[b_ptr], %[b_ptr], #48\n"
290                 "add	%[a_ptr], %[a_ptr], #32\n"
291                 "str	q24, [%[c_ptr], #32]\n"
292                 "fmla	v25.4s, %[b2].4s, %[a0].s[1]\n"
293                 "str	q9, [%[c_ptr], #48]\n"
294 
295                 "fmla	v10.4s, %[b0].4s, %[a0].s[2]\n"
296                 "str	q17, [%[c_ptr], #64]\n"
297                 "fmla	v18.4s, %[b1].4s, %[a0].s[2]\n"
298                 "str	q25, [%[c_ptr], #80]\n"
299                 "fmla	v26.4s, %[b2].4s, %[a0].s[2]\n"
300                 "str	q10, [%[c_ptr], #96]\n"
301 
302                 "fmla	v11.4s, %[b0].4s, %[a0].s[3]\n"
303                 "str	q18, [%[c_ptr], #112]\n"
304                 "fmla	v19.4s, %[b1].4s, %[a0].s[3]\n"
305                 "str	q26, [%[c_ptr], #128]\n"
306                 "fmla	v27.4s, %[b2].4s, %[a0].s[3]\n"
307                 "str	q11, [%[c_ptr], #144]\n"
308 
309                 "fmla 	v12.4s, %[b0].4s, %[a1].s[0]\n"
310                 "str	q19, [%[c_ptr], #160]\n"
311                 "fmla	v20.4s, %[b1].4s, %[a1].s[0]\n"
312                 "str	q27, [%[c_ptr], #176]\n"
313                 "fmla	v28.4s, %[b2].4s, %[a1].s[0]\n"
314                 "str	q12, [%[c_ptr], #192]\n"
315 
316                 "fmla   v13.4s, %[b0].4s, %[a1].s[1]\n"
317                 "str	q20, [%[c_ptr], #208]\n"
318                 "fmla	v21.4s, %[b1].4s, %[a1].s[1]\n"
319                 "str	q28, [%[c_ptr], #224]\n"
320                 "fmla	v29.4s, %[b2].4s, %[a1].s[1]\n"
321                 "str	q13, [%[c_ptr], #240]\n"
322 
323                 "fmla	v14.4s, %[b0].4s, %[a1].s[2]\n"
324                 "str	q21, [%[c_ptr], #256]\n"
325                 "fmla	v22.4s, %[b1].4s, %[a1].s[2]\n"
326                 "str	q29, [%[c_ptr], #272]\n"
327                 "fmla	v30.4s, %[b2].4s, %[a1].s[2]\n"
328                 "str	q14, [%[c_ptr], #288]\n"
329 
330                 "fmla	v15.4s, %[b0].4s, %[a1].s[3]\n"
331                 "str	q22, [%[c_ptr], #304]\n"
332                 "fmla	v23.4s, %[b1].4s, %[a1].s[3]\n"
333                 "str	q30, [%[c_ptr], #320]\n"
334                 "fmla	v31.4s, %[b2].4s, %[a1].s[3]\n"
335                 "str	q15, [%[c_ptr], #336]\n"
336 
337                 // Common tail
338                 "3:\n"
339                 "str	q23, [%[c_ptr], #352]\n"
340                 "str	q31, [%[c_ptr], #368]\n"
341                 "add	%[c_ptr], %[c_ptr], #384\n"
342             :
343               [a_ptr] "+r" (a_ptr), [b_ptr] "+r" (b_ptr), [c_ptr] "+r" (c_ptr),
344               [a0] "+w" (a0), [a1] "+w" (a1), [a0a] "+w" (a0a), [a1a] "+w" (a1a),
345               [b0] "+w" (b0), [b1] "+w" (b1), [b2] "+w" (b2), [k] "+r" (k)
346             : [oddk] "r" (oddk)
347             : "x20", "x21", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18",
348               "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "cc"
349             );
350         }
351     }
352 }
353 
354 } // namespace arm_gemm
355 
356 #endif
357