xref: /aosp_15_r20/external/XNNPACK/src/f32-gemm/4x8-aarch32-neon-cortex-a53.S.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
1// Copyright 2019 Google LLC
2//
3// This source code is licensed under the BSD-style license found in the
4// LICENSE file in the root directory of this source tree.
5
6#include <xnnpack/assembly.h>
7
8.syntax unified
9
10# LINT.IfChange
11// void xnn_f32_gemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_cortex_a53(
12//     size_t mr,                            r0
13//     size_t nc,                            r1
14//     size_t kc,                            r2 -> r5 -> sp + 0
15//     const uint8_t*restrict a,             r3
16//     size_t a_stride,          sp + 100 -> (r7)
17//     const void*restrict w,    sp + 104 -> r9
18//     uint8_t*restrict c,       sp + 108 -> r11
19//     size_t cm_stride,         sp + 112 -> (r6)
20//     size_t cn_stride,         sp + 116 -> (r0)
21//     const union xnn_f32_minmax_params params)  sp + 120 -> (r5)
22
23// d8-d15, r4-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.
24
25// Register usage
26
27// r0, r2   scratch temporaries for loads
28// r14 (lr) unused
29
30// A0   r3  d0
31// A1  r12  d1
32// A2  r10  d2
33// A3   r7  d3
34
35// B    r9  d8,  d9, d10, d11
36// B       d12, d13, d14, d15
37
38// C0  r11 d16-d17  q8  d18-d19  q9
39// C1   r4 d20-d21 q10  d22-d23 q11
40// C2   r8 d24-d25 q12  d26-d27 q13
41// C3   r6 d28-d29 q14  d30-d31 q15
42
43// Clamp (r5) d4 d5 d6 d7
44
45BEGIN_FUNCTION xnn_f32_gemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_cortex_a53
46        .arm
47#ifndef __APPLE__
48        .arch   armv7-a
49        .fpu    neon
50#endif
51        # Push 100 bytes
52        # r2 will be reloaded in outer loop
53        VPUSH   {d8-d15}                                // 64
54        PUSH    {r2, r4, r5, r6, r7, r8, r9, r10, r11}  // +36 = 100
55
56        LDR     r7, [sp, 100]           // a_stride
57        LDR     r11, [sp, 108]          // c
58        LDR     r6, [sp, 112]           // cm_stride
59        LDR     r9, [sp, 104]           // w
60
61        # Clamp A and C pointers
62        CMP     r0, 2                   // if mr >= 2
63        ADD     r12, r3, r7             //   a1 = a0 + a_stride
64        ADD     r4, r11, r6             //   c1 = c0 + cm_stride
65        MOVLO   r12, r3                 // a1
66        MOVLO   r4, r11                 // c1
67                                        // if mr > 2
68        ADD     r10, r12, r7            //   a2 = a1 + a_stride
69        ADD     r8, r4, r6              //   c2 = c1 + cm_stride
70        MOVLS   r10, r12                // a2
71        MOVLS   r8, r4                  // c2
72
73        CMP     r0, 4                   // if mr >=4
74        ADD     r7, r10, r7             //   a3 = a2 + a_stride
75        ADD     r6, r8, r6              //   c3 = c2 + cm_stride
76        MOVLO   r7, r10                 // a3
77        MOVLO   r6, r8                  // c3
78
79        .p2align 3
800:
81        # Load initial bias from w into accumulators
82        VLDM    r9!, {d16-d19}          // Bias
83
84        SUBS    r5, r2, 16              // kc - 16
85        $if PREFETCH:
86          PLD     [r3,  0]                // Prefetch A
87        $if PREFETCH:
88          PLD     [r3, 64]
89        VMOV    q10, q8
90        $if PREFETCH:
91          PLD     [r12,  0]
92        $if PREFETCH:
93          PLD     [r12, 64]
94        VMOV    q11, q9
95        $if PREFETCH:
96          PLD     [r10,  0]
97        $if PREFETCH:
98          PLD     [r10, 64]
99        VMOV    q12, q8
100        $if PREFETCH:
101          PLD     [r7,  0]
102        $if PREFETCH:
103          PLD     [r7, 64]
104        VMOV    q13, q9
105        $if PREFETCH:
106          PLD     [r9,   0]               // Prefetch B
107        $if PREFETCH:
108          PLD     [r9,  64]
109        VMOV    q14, q8
110        $if PREFETCH:
111          PLD     [r9, 128]
112        $if PREFETCH:
113          PLD     [r9, 192]
114        VMOV    q15, q9
115        $if PREFETCH:
116          PLD     [r9, 256]
117        $if PREFETCH:
118          PLD     [r9, 320]
119        BLO     4f                      // less than 4 channels?
120
121        # Prologue
122        VLD1.32 {d0},  [r3]!            // A0
123        VLD1.32 {d1}, [r12]!            // A1
124        VLD1.32 {d2}, [r10]!            // A2
125        VLD1.32 {d3},  [r7]!            // A3
126        SUBS    r5, r5, 16
127        VLDM    r9, {d8-d11}            // B0
128        LDR     r0, [r9, 56]            // B1 low   VMOV is in BLOCK 0
129        LDR     r2, [r9, 60]            // B1 high
130        VLDR    d13, [r9, 40]           // B1
131
132        BLO     2f                      // less than 4 channels?  skip main loop
133
134        # Main loop - 4 floats of A (16 bytes)
135        # 32 FMA + 8 LD64 A + 8 LDR B
136        .p2align 3
1371:
138        # First group of 16 FMA, Second group loads
139        # BLOCK 0
140        VLD1.32 {d4}, [r3]!             // A0
141        VMOV    d15, r0, r2             // b1 VMOV b from second group
142        VMLA.F32 q8, q4, d0[0]
143        LDR     r0, [r12]               // A1 low
144        VMLA.F32 q10, q4, d1[0]
145        LDR     r2, [r12, 4]            // A1 high
146        VMLA.F32 q12, q4, d2[0]
147        $if PREFETCH:
148          PLD     [r3, 128]               // Prefetch A0
149
150        # BLOCK 1
151        VLDR    d12, [r9, 32]           // B1
152        VMOV    d5, r0, r2              // a1 VMOV
153        VMLA.F32 q14, q4, d3[0]
154        LDR     r0, [r9, 72]            // B0 low
155        VMLA.F32 q9, q5, d0[0]
156        LDR     r2, [r9, 76]            // B0 high
157        VMLA.F32 q11, q5, d1[0]
158        $if PREFETCH:
159          PLD     [r12, 128]              // Prefetch A1
160
161        # BLOCK 2
162        VLD1.32 {d6}, [r10]!            // A2
163        VMOV    d9, r0, r2              // b0 VMOV
164        VMLA.F32 q13, q5, d2[0]
165        LDR     r0, [r7]                // A3 low
166        VMLA.F32 q15, q5, d3[0]
167        LDR     r2, [r7, 4]             // A3 high
168        VMLA.F32 q8, q6, d0[1]
169        $if PREFETCH:
170          PLD     [r10, 128]              // Prefetch A2
171
172        # BLOCK 3
173        VLDR    d14, [r9, 48]           // B1
174        VMOV    d7, r0, r2              // a3 VMOV
175        VMLA.F32 q10, q6, d1[1]
176        LDR     r0, [r9, 88]            // B0 low
177        VMLA.F32 q12, q6, d2[1]
178        LDR     r2, [r9, 92]            // B0 high
179        VMLA.F32 q14, q6, d3[1]
180        $if PREFETCH:
181          PLD     [r7, 128]               // Prefetch A3
182
183        # BLOCK 4
184        VLDR    d8, [r9, 64]            // B0
185        VMOV    d11, r0, r2             // B0 VMOV
186        VMLA.F32 q9, q7, d0[1]
187        LDR     r0, [r9, 104]           // B1 low   VMOV is in BLOCK 0
188        VMLA.F32 q11, q7, d1[1]
189        LDR     r2, [r9, 108]           // B1 high
190        VMLA.F32 q13, q7, d2[1]
191        $if PREFETCH:
192          PLD     [r9, 384]               // Prefetch B
193
194        # BLOCK 5
195        VLDR    d10, [r9, 80]           // B0
196        VMOV    d13, r0, r2             // b1 VMOV b from second group
197        VMLA.F32 q15, q7, d3[1]
198        LDR     r0, [r9, 120]           // B1 low   VMOV is in BLOCK 0
199        NOP
200        LDR     r2, [r9, 124]           // B1 high
201        NOP
202        $if PREFETCH:
203          PLD     [r9, 448]               // Prefetch B
204
205        # Second group of 16 FMA, First group of loads
206        # BLOCK 0
207        VLD1.32 {d0}, [r3]!             // A0
208        VMOV    d15, r0, r2             // b1 VMOV b from second group
209        VMLA.F32 q8, q4, d4[0]
210        LDR     r0, [r12, 8]            // A1 low
211        VMLA.F32 q10, q4, d5[0]
212        LDR     r2, [r12, 12]           // A1 high
213        VMLA.F32 q12, q4, d6[0]
214        # NOP
215
216        # BLOCK 1
217        VLDR    d12, [r9, 96]           // B1
218        VMOV    d1, r0, r2              // a1 VMOV
219        VMLA.F32 q14, q4, d7[0]
220        LDR     r0, [r9, 136]           // B0 low
221        VMLA.F32 q9, q5, d4[0]
222        LDR     r2, [r9, 140]           // B0 high
223        VMLA.F32 q11, q5, d5[0]
224        # NOP
225
226        # BLOCK 2
227        VLD1.32 {d2}, [r10]!            // A2
228        VMOV    d9, r0, r2              // b0 VMOV
229        VMLA.F32 q13, q5, d6[0]
230        LDR     r0, [r7, 8]             // A3 low
231        VMLA.F32 q15, q5, d7[0]
232        LDR     r2, [r7, 12]            // A3 high
233        VMLA.F32 q8, q6, d4[1]
234        # NOP
235
236        # BLOCK 3
237        VLDR    d14, [r9, 112]          // B1
238        VMOV    d3, r0, r2              // a3 VMOV
239        VMLA.F32 q10, q6, d5[1]
240        LDR     r0, [r9, 152]           // B0 low
241        VMLA.F32 q12, q6, d6[1]
242        LDR     r2, [r9, 156]           // B0 high
243        VMLA.F32 q14, q6, d7[1]
244        ADD     r12, r12, 16            // A1++
245
246        # BLOCK 4
247        VLDR    d8, [r9, 128]           // B0
248        VMOV    d11, r0, r2             // B0 VMOV
249        VMLA.F32 q9, q7, d4[1]
250        LDR     r0, [r9, 168]           // B1 low
251        VMLA.F32 q11, q7, d5[1]
252        LDR     r2, [r9, 172]           // B1 high
253        VMLA.F32 q13, q7, d6[1]
254        ADD     r7, r7, 16              // A3++
255
256        # BLOCK 5
257        VLDR    d10, [r9, 144]          // B0
258        VMOV    d13, r0, r2             // b1 VMOV b
259        VMLA.F32 q15, q7, d7[1]
260        LDR     r0, [r9, 184]           // B1 low   VMOV is in BLOCK 0
261        SUBS    r5, r5, 16
262        LDR     r2, [r9, 188]           // B1 high
263        ADD     r9, r9, 128             // B++
264        BHS     1b
265
266        # Epilogue - 4 floats of A (16 bytes)
2672:
268        # First group of 16 FMA, Second group loads
269        # BLOCK 0
270        VLD1.32 {d4}, [r3]!             // A0
271        VMOV    d15, r0, r2             // b1 VMOV b from second group
272        VMLA.F32 q8, q4, d0[0]
273        LDR     r0, [r12]               // A1 low
274        VMLA.F32 q10, q4, d1[0]
275        LDR     r2, [r12, 4]            // A1 high
276        VMLA.F32 q12, q4, d2[0]
277        # NOP
278
279        # BLOCK 1
280        VLDR    d12, [r9, 32]           // B1
281        VMOV    d5, r0, r2              // a1 VMOV
282        VMLA.F32 q14, q4, d3[0]
283        LDR     r0, [r9, 72]            // B0 low
284        VMLA.F32 q9, q5, d0[0]
285        LDR     r2, [r9, 76]            // B0 high
286        VMLA.F32 q11, q5, d1[0]
287        # NOP
288
289        # BLOCK 2
290        VLD1.32 {d6}, [r10]!            // A2
291        VMOV    d9, r0, r2              // b0 VMOV
292        VMLA.F32 q13, q5, d2[0]
293        LDR     r0, [r7]                // A3 low
294        VMLA.F32 q15, q5, d3[0]
295        LDR     r2, [r7, 4]             // A3 high
296        VMLA.F32 q8, q6, d0[1]
297        # NOP
298
299        # BLOCK 3
300        VLDR    d14, [r9, 48]           // B1
301        VMOV    d7, r0, r2              // a3 VMOV
302        VMLA.F32 q10, q6, d1[1]
303        LDR     r0, [r9, 88]            // B0 low
304        VMLA.F32 q12, q6, d2[1]
305        LDR     r2, [r9, 92]            // B0 high
306        VMLA.F32 q14, q6, d3[1]
307        # NOP
308
309        # BLOCK 4
310        VLDR    d8, [r9, 64]            // B0
311        VMOV    d11, r0, r2             // B0 VMOV
312        VMLA.F32 q9, q7, d0[1]
313        LDR     r0, [r9, 104]           // B1 low
314        VMLA.F32 q11, q7, d1[1]
315        LDR     r2, [r9, 108]           // B1 high
316        VMLA.F32 q13, q7, d2[1]
317        # NOP
318
319        # BLOCK 5
320        VLDR    d10, [r9, 80]           // B0
321        VMOV    d13, r0, r2             // b1 VMOV b
322        VMLA.F32 q15, q7, d3[1]
323        LDR     r0, [r9, 120]           // B1 low   VMOV is in BLOCK 0
324        NOP
325        LDR     r2, [r9, 124]           // B1 high
326        NOP
327        NOP
328
329        # Second group of 16 FMA, First group of loads
330        # BLOCK 0
331        VLDR    d12, [r9, 96]           // B1
332        VMOV    d15, r0, r2             // b1 VMOV b from second group
333        VMLA.F32 q8, q4, d4[0]
334        VMLA.F32 q10, q4, d5[0]
335        VMLA.F32 q12, q4, d6[0]
336
337        # BLOCK 1
338        VLDR    d14, [r9, 112]          // B1
339        VMLA.F32 q14, q4, d7[0]
340        VMLA.F32 q9, q5, d4[0]
341        VMLA.F32 q11, q5, d5[0]
342        ADD     r12, r12, 8             // A1++
343
344        # BLOCK 2
345        ADD     r7, r7, 8               // A3++ VLDR B1 lands here
346        ADD     r9, r9, 128             // B++
347        VMLA.F32 q13, q5, d6[0]
348        VMLA.F32 q15, q5, d7[0]
349        VMLA.F32 q8, q6, d4[1]
350
351        # BLOCK 3
352        VMLA.F32 q10, q6, d5[1]
353        VMLA.F32 q12, q6, d6[1]
354        VMLA.F32 q14, q6, d7[1]
355        TST     r5, 15
356
357        # BLOCK 4
358        VMLA.F32 q9, q7, d4[1]
359        VMLA.F32 q11, q7, d5[1]
360        VMLA.F32 q13, q7, d6[1]
361
362        # BLOCK 5
363        VMLA.F32 q15, q7, d7[1]
364
365        # Is there a remainder?- 1 to 3 floats of A (4, 8 or 12 bytes)
366        BNE     4f
367
368        .p2align 3
3693:
370        # Load params pointer
371        LDR     r0, [sp, 116]           // cn_stride
372        LDR     r5, [sp, 120]           // params
373        LDR     r2, [sp]                // kc
374        SUBS    r1, r1, 8
375
376        # Load min/max values
377        VLD1.32 {d4[],d5[]}, [r5]!
378        VLD1.32 {d6[],d7[]}, [r5]
379
380        # Clamp
381        VMAX.F32 q8,  q8, q2
382        VMAX.F32 q9,  q9, q2
383        VMAX.F32 q10, q10, q2
384        VMAX.F32 q11, q11, q2
385        VMAX.F32 q12, q12, q2
386        VMAX.F32 q13, q13, q2
387        VMAX.F32 q14, q14, q2
388        VMAX.F32 q15, q15, q2
389        VMIN.F32 q8,  q8, q3
390        VMIN.F32 q9,  q9, q3
391        VMIN.F32 q10, q10, q3
392        VMIN.F32 q11, q11, q3
393        VMIN.F32 q12, q12, q3
394        VMIN.F32 q13, q13, q3
395        VMIN.F32 q14, q14, q3
396        VMIN.F32 q15, q15, q3
397
398        # Store full 4 x 8
399        BLO     6f
400        VST1.32 {d16-d19}, [r11], r0
401        SUB     r7, r7, r2
402        VST1.32 {d20-d23}, [r4], r0
403        SUB     r10, r10, r2
404        VST1.32 {d24-d27}, [r8], r0
405        SUB     r12, r12, r2
406        VST1.32 {d28-d31}, [r6], r0
407        SUB     r3, r3, r2
408        BHI     0b
409
410        ADD     sp, sp, 4
411        POP     {r4, r5, r6, r7, r8, r9, r10, r11}
412        VPOP    {d8-d15}
413        BX      lr
414
415        .p2align 3
4164:
417        # Is there a remainder?- 2 floats of A (8 bytes)
418        TST     r5, 8
419        BEQ     5f
420
421        # Remainder - 2 floats of A (8 bytes)
422        VLD1.32 {d0}, [r3]!             // A0
423        VLDM    r9!, {d8-d11}           // B0
424        VLD1.32 {d1}, [r12]!            // A1
425        VLD1.32 {d2}, [r10]!            // A2
426        VLD1.32 {d3}, [ r7]!            // A3
427
428        VMLA.F32 q8, q4, d0[0]
429        VMLA.F32 q9, q5, d0[0]
430        VMLA.F32 q10, q4, d1[0]
431        VMLA.F32 q11, q5, d1[0]
432        VLDM    r9!, {d12-d15}          // B1
433        VMLA.F32 q12, q4, d2[0]
434        VMLA.F32 q13, q5, d2[0]
435        VMLA.F32 q14, q4, d3[0]
436        VMLA.F32 q15, q5, d3[0]
437        VMLA.F32 q8, q6, d0[1]
438        VMLA.F32 q9, q7, d0[1]
439        VMLA.F32 q10, q6, d1[1]
440        VMLA.F32 q11, q7, d1[1]
441        VMLA.F32 q12, q6, d2[1]
442        VMLA.F32 q13, q7, d2[1]
443        VMLA.F32 q14, q6, d3[1]
444        VMLA.F32 q15, q7, d3[1]
445
446        # Is there a remainder?- 1 float of A (4 bytes)
447        TST     r5, 4
448        BEQ     3b
449
4505:
451        # Remainder- 1 float of A (4 bytes)
452        VLDM    r3!,  {s0}              // A0
453        VLDM    r9!, {d8-d11}           // B0
454        VLDM    r12!, {s2}              // A1
455        VLDM    r10!, {s4}              // A2
456        VLDM    r7!, {s6}               // A3
457        VMLA.F32 q8, q4, d0[0]
458        VMLA.F32 q9, q5, d0[0]
459        VMLA.F32 q10, q4, d1[0]
460        VMLA.F32 q11, q5, d1[0]
461        VMLA.F32 q12, q4, d2[0]
462        VMLA.F32 q13, q5, d2[0]
463        VMLA.F32 q14, q4, d3[0]
464        VMLA.F32 q15, q5, d3[0]
465        B       3b
466
467        # Store odd width
4686:
469        TST     r1, 4
470        BEQ     7f
471        VST1.32 {d16-d17}, [r11]!
472        VST1.32 {d20-d21},  [r4]!
473        VMOV    q8,  q9
474        VMOV    q10, q11
475        VST1.32 {d24-d25},  [r8]!
476        VST1.32 {d28-d29},  [r6]!
477        VMOV    q12, q13
478        VMOV    q14, q15
479
4807:
481        TST     r1, 2
482        BEQ     8f
483        VST1.32 {d16}, [r11]!
484        VST1.32 {d20},  [r4]!
485        VMOV    d16, d17
486        VMOV    d20, d21
487        VST1.32 {d24},  [r8]!
488        VST1.32 {d28},  [r6]!
489        VMOV    d24, d25
490        VMOV    d28, d29
491
4928:
493        TST     r1, 1
494        BEQ     9f
495        VST1.32 {d16[0]}, [r11]
496        VST1.32 {d20[0]},  [r4]
497        VST1.32 {d24[0]},  [r8]
498        VST1.32 {d28[0]},  [r6]
499
5009:
501        ADD     sp, sp, 4
502        POP     {r4, r5, r6, r7, r8, r9, r10, r11}
503        VPOP    {d8-d15}
504        BX      lr
505
506END_FUNCTION xnn_f32_gemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_cortex_a53
507# LINT.ThenChange(4x8-aarch32-neon-cortex-a53.cc)
508
509#ifdef __ELF__
510.section ".note.GNU-stack","",%progbits
511#endif
512