xref: /aosp_15_r20/external/XNNPACK/src/f32-igemm/4x8-minmax-aarch32-neon-cortex-a55.S (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_igemm_minmax_ukernel_4x8__aarch32_neon_cortex_a55(
12//     size_t mr,                            r0
13//     size_t nc,                            r1
14//     size_t kc,                            r2 -> r5
15//     size_t ks,                            r3 -> sp + 64 -> r14
16//     const float**restrict a,  sp + 104 -> (r5)
17//     const void*restrict w,    sp + 108 -> r9
18//     uint8_t*restrict c,       sp + 112 -> r11
19//     size_t cm_stride,         sp + 116 -> (r6)
20//     size_t cn_stride,         sp + 120 -> (r0)
21//     size_t a_offset,          sp + 124 -> (r5)
22//     const float* zero,        sp + 128 -> (r0)
23//     minmax_params*params,     sp + 132 -> (r5)
24
25// d8-d15, r4-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.
26
27// Register usage
28
29// A0   r3  d0
30// A1  r12  d1
31// A2  r10  d2
32// A3   r7  d3
33
34// B    r9  d8,  d9, d10, d11
35// B       d12, d13, d14, d15
36
37// C0  r11 d16-d17  q8  d18-d19  q9
38// C1   r4 d20-d21 q10  d22-d23 q11
39// C2   r8 d24-d25 q12  d26-d27 q13
40// C3   r6 d28-d29 q14  d30-d31 q15
41
42// Clamp (r5) d4 d5 d6 d7
43
44BEGIN_FUNCTION xnn_f32_igemm_minmax_ukernel_4x8__aarch32_neon_cortex_a55
45        .arm
46#ifndef __APPLE__
47        .arch   armv7-a
48        .fpu    neon
49#endif
50        # Push 104 bytes
51        PUSH    {r3, r4, r5, r6, r7, r8, r9, r10, r11, lr}   // +40
52        VPUSH   {d8-d15}                                     // +64 = 104
53
54        LDR     r11, [sp, 112]          // c
55        LDR     r6, [sp, 116]           // cm_stride
56        LDR     r5, [sp, 104]           // a
57        LDR     r9, [sp, 108]           // w
58        MOV     r14, r3                 // p = ks
59
60        # Clamp C pointers
61        CMP     r0, 2                   // if mr >= 2
62        ADD     r4, r11, r6             //   c1 = c0 + cm_stride
63        MOVLO   r4, r11                 // c1
64                                     // if mr > 2
65        ADD     r8, r4, r6              //   c2 = c1 + cm_stride
66        MOVLS   r8, r4                  // c2
67        CMP     r0, 4                   // if mr >=4
68        ADD     r6, r8, r6              //   c3 = c2 + cm_stride
69        MOVLO   r6, r8                  // c3
70
71
72        .p2align 3
730:
74        # Load initial bias from w into accumulators
75        VLDM    r9!, {d16-d19}          // Bias
76
77        VMOV    q10, q8
78        VMOV    q11, q9
79        VMOV    q12, q8
80        VMOV    q13, q9
81        PLD     [r9,   0]               // Prefetch B
82        PLD     [r9,  64]
83        VMOV    q14, q8
84        PLD     [r9, 128]
85        PLD     [r9, 192]
86        VMOV    q15, q9
87        PLD     [r9, 256]
88        PLD     [r9, 320]
89
901:
91        # Load next 4 A pointers
92        LDR     r3, [r5,  0]
93        LDR     r12, [r5,  4]
94        LDR     r10, [r5,  8]
95        LDR     r7, [r5, 12]
96        ADD     r5, r5, 16
97        PLD     [r3,  0]                // Prefetch A
98        STR     r5, [sp, 104]           // a
99        PLD     [r3, 64]
100        LDR     r0, [sp, 128]           // zero
101        PLD     [r12,  0]
102        LDR     r5, [sp, 124]           // a_offset
103        PLD     [r12, 64]
104        PLD     [r10,  0]
105        PLD     [r10, 64]
106        PLD     [r7,  0]
107        PLD     [r7, 64]
108
109        # Add a_offset
110        CMP     r3,  r0                 // if a0 == zero
111        ADD     r3,  r3, r5             // a0 += a_offset
112        MOVEQ   r3,  r0                 //   a0 = zero, else += a0 + a_offset
113        CMP     r12,  r0                // if a1 == zero
114        ADD     r12, r12, r5            // a1 += a_offset
115        MOVEQ   r12,  r0                //   a1 = zero, else += a1 + a_offset
116        CMP     r10,  r0                // if a2 == zero
117        ADD     r10, r10, r5            // a2 += a_offset
118        MOVEQ   r10,  r0                //   a2 = zero, else += a2 + a_offset
119        CMP     r7,  r0                 // if a3 == zero
120        ADD     r7,  r7, r5             // a3 += a_offset
121        MOVEQ   r7,  r0                 //   a3 = zero, else += a3 + a_offset
122
123        SUBS    r5, r2, 16              // kc - 16
124        BLO     5f                      // less than 4 channels?
125
126        # Prologue
127        VLD1.32 {d0},  [r3]!            // A0
128        VLD1.32 {d1}, [r12]!            // A1
129        VLD1.32 {d2}, [r10]!            // A2
130        VLD1.32 {d3},  [r7]!            // A3
131        SUBS    r5, r5, 16
132        VLDM    r9, {d8-d11}            // B0
133        VLDR    d15, [r9, 56]           // B1CK 0
134        VLDR    d13, [r9, 40]           // B1
135
136        BLO     3f                      // less than 4 channels?  skip main loop
137
138        # Main loop - 4 floats of A (16 bytes)
139        # 32 FMA + 8 LD64 A + 8 LDR B
140        .p2align 3
1412:
142        # First group of 16 FMA, Second group loads
143        # BLOCK 0
144        VMLA.F32 q8, q4, d0[0]
145        VLD1.32 {d4}, [r3]!             // A0
146        VMLA.F32 q10, q4, d1[0]
147        VLD1.32 {d5}, [r12]!            // A1
148        VMLA.F32 q12, q4, d2[0]
149
150        # BLOCK 1
151        VMLA.F32 q14, q4, d3[0]
152        VLDR    d12, [r9, 32]           // B1
153        VMLA.F32 q9, q5, d0[0]
154        VLDR    d9, [r9, 72]            // B0
155        VMLA.F32 q11, q5, d1[0]
156
157        # BLOCK 2
158        VMLA.F32 q13, q5, d2[0]
159        VLD1.32 {d6}, [r10]!            // A2
160        VMLA.F32 q15, q5, d3[0]
161        VLD1.32 {d7}, [r7]!             // A3
162        VMLA.F32 q8, q6, d0[1]
163
164        # BLOCK 3
165        VMLA.F32 q10, q6, d1[1]
166        VLDR    d14, [r9, 48]           // B1
167        VMLA.F32 q12, q6, d2[1]
168        VLDR    d11, [r9, 88]           // B0
169        VMLA.F32 q14, q6, d3[1]
170
171        # BLOCK 4
172        VMLA.F32 q9, q7, d0[1]
173        VLDR    d8, [r9, 64]            // B0
174        VMLA.F32 q11, q7, d1[1]
175        VLDR    d13, [r9, 104]          // B1
176        VMLA.F32 q13, q7, d2[1]
177        VLDR    d10, [r9, 80]           // B0
178
179        # BLOCK 5
180        VMLA.F32 q15, q7, d3[1]
181        VLDR    d15, [r9, 120]          // B1
182
183        # Second group of 16 FMA, First group of loads
184        # BLOCK 0
185        VMLA.F32 q8, q4, d4[0]
186        VLD1.32 {d0}, [r3]!             // A0
187        VMLA.F32 q10, q4, d5[0]
188        VLD1.32 {d1}, [r12]!            // A1
189        VMLA.F32 q12, q4, d6[0]
190
191        # BLOCK 1
192        VMLA.F32 q14, q4, d7[0]
193        VLDR    d12, [r9, 96]           // B1
194        VMLA.F32 q9, q5, d4[0]
195        VLDR    d9, [r9, 136]           // B0
196        VMLA.F32 q11, q5, d5[0]
197
198        # BLOCK 2
199        VMLA.F32 q13, q5, d6[0]
200        VLD1.32 {d2}, [r10]!            // A2
201        VMLA.F32 q15, q5, d7[0]
202        VLD1.32 {d3}, [r7]!             // A3
203        VMLA.F32 q8, q6, d4[1]
204        SUBS    r5, r5, 16
205
206        # BLOCK 3
207        VMLA.F32 q10, q6, d5[1]
208        VLDR    d14, [r9, 112]          // B1
209        VMLA.F32 q12, q6, d6[1]
210        VLDR    d11, [r9, 152]          // B0
211        VMLA.F32 q14, q6, d7[1]
212
213        # BLOCK 4
214        VMLA.F32 q9, q7, d4[1]
215        VLDR    d8, [r9, 128]           // B0
216        VMLA.F32 q11, q7, d5[1]
217        VLDR    d13, [r9, 168]          // B1
218        VMLA.F32 q13, q7, d6[1]
219        VLDR    d10, [r9, 144]          // B0
220
221        # BLOCK 5
222        VMLA.F32 q15, q7, d7[1]
223        VLDR    d15, [r9, 184]          // B1
224        ADD     r9, r9, 128             // B++
225        BHS     2b
226
227        # Epilogue - 4 floats of A (16 bytes)
2283:
229        # First group of 16 FMA, Second group loads
230        # BLOCK 0
231        VMLA.F32 q8, q4, d0[0]
232        VLD1.32 {d4}, [r3]!             // A0
233        VMLA.F32 q10, q4, d1[0]
234        VLD1.32 {d5}, [r12]!            // A1
235        VMLA.F32 q12, q4, d2[0]
236
237        # BLOCK 1
238        VMLA.F32 q14, q4, d3[0]
239        VLDR    d12, [r9, 32]           // B1
240        VMLA.F32 q9, q5, d0[0]
241        VLDR    d9, [r9, 72]            // B0
242        VMLA.F32 q11, q5, d1[0]
243
244        # BLOCK 2
245        VMLA.F32 q13, q5, d2[0]
246        VLD1.32 {d6}, [r10]!            // A2
247        VMLA.F32 q15, q5, d3[0]
248        VLD1.32 {d7}, [r7]!             // A3
249        VMLA.F32 q8, q6, d0[1]
250
251        # BLOCK 3
252        VMLA.F32 q10, q6, d1[1]
253        VLDR    d14, [r9, 48]           // B1
254        VMLA.F32 q12, q6, d2[1]
255        VLDR    d11, [r9, 88]           // B0
256        VMLA.F32 q14, q6, d3[1]
257
258        # BLOCK 4
259        VMLA.F32 q9, q7, d0[1]
260        VLDR    d8, [r9, 64]            // B0
261        VMLA.F32 q11, q7, d1[1]
262        VLDR    d13, [r9, 104]          // B1
263        VMLA.F32 q13, q7, d2[1]
264        VLDR    d10, [r9, 80]           // B0
265
266        # BLOCK 5
267        VMLA.F32 q15, q7, d3[1]
268        VLDR    d15, [r9, 120]          // B1
269
270        # Second group of 16 FMA, First group of loads
271        # BLOCK 0
272        VMLA.F32 q8, q4, d4[0]
273        VLDR    d12, [r9, 96]           // B1
274        VMLA.F32 q10, q4, d5[0]
275        VMLA.F32 q12, q4, d6[0]
276
277        # BLOCK 1
278        VMLA.F32 q14, q4, d7[0]
279        VLDR    d14, [r9, 112]          // B1
280        VMLA.F32 q9, q5, d4[0]
281        VMLA.F32 q11, q5, d5[0]
282
283        # BLOCK 2
284        VMLA.F32 q13, q5, d6[0]
285        VMLA.F32 q15, q5, d7[0]
286        VMLA.F32 q8, q6, d4[1]
287        ADD     r9, r9, 128             // B++
288
289        # BLOCK 3
290        VMLA.F32 q10, q6, d5[1]
291        VMLA.F32 q12, q6, d6[1]
292        VMLA.F32 q14, q6, d7[1]
293        TST     r5, 15
294
295        # BLOCK 4
296        VMLA.F32 q9, q7, d4[1]
297        VMLA.F32 q11, q7, d5[1]
298        VMLA.F32 q13, q7, d6[1]
299
300        # BLOCK 5
301        VMLA.F32 q15, q7, d7[1]
302
303        # Is there a remainder?- 1 to 3 floats of A (4, 8 or 12 bytes)
304        BNE     5f
305
306        .p2align 3
3074:
308        LDR     r5, [sp, 104]           // a
309        SUBS    r14, r14, 16            // ks -= MR * sizeof(void*)
310
311        # ks loop
312        BHI     1b
313
314        # Load params pointer
315        LDR     r0, [sp, 132]           // params
316        LDR     r14, [sp, 64]           // p = ks
317        # Load min/max values
318        VLD1.32 {d4[],d5[]}, [r0]!
319        VLD1.32 {d6[],d7[]}, [r0]
320        SUBS    r1, r1, 8
321        LDR     r0, [sp, 120]           // cn_stride
322
323        # Clamp
324        VMAX.F32 q8,  q8, q2
325        VMAX.F32 q9,  q9, q2
326        VMAX.F32 q10, q10, q2
327        VMAX.F32 q11, q11, q2
328        VMAX.F32 q12, q12, q2
329        VMAX.F32 q13, q13, q2
330        VMAX.F32 q14, q14, q2
331        VMAX.F32 q15, q15, q2
332        VMIN.F32 q8,  q8, q3
333        VMIN.F32 q9,  q9, q3
334        VMIN.F32 q10, q10, q3
335        VMIN.F32 q11, q11, q3
336        VMIN.F32 q12, q12, q3
337        VMIN.F32 q13, q13, q3
338        VMIN.F32 q14, q14, q3
339        VMIN.F32 q15, q15, q3
340
341        # Store full 4 x 8
342        BLO     7f
343        VST1.32 {d28-d31},  [r6], r0
344        VST1.32 {d24-d27},  [r8], r0
345        VST1.32 {d20-d23},  [r4], r0
346        VST1.32 {d16-d19}, [r11], r0
347
348        SUB     r5, r5, r14             // a -= ks
349
350        BHI     0b
351
352        VPOP    {d8-d15}
353        ADD     sp, sp, 4               // skip r3
354        POP     {r4, r5, r6, r7, r8, r9, r10, r11, pc}
355
356        .p2align 3
3575:
358        # Is there a remainder?- 2 floats of A (8 bytes)
359        TST     r5, 8
360        BEQ     6f
361
362        # Remainder - 2 floats of A (8 bytes)
363        VLD1.32 {d0}, [r3]!             // A0
364        VLDM    r9!, {d8-d11}           // B0
365        VLD1.32 {d1}, [r12]!            // A1
366        VLD1.32 {d2}, [r10]!            // A2
367        VLD1.32 {d3}, [ r7]!            // A3
368
369        VMLA.F32 q8, q4, d0[0]
370        VMLA.F32 q9, q5, d0[0]
371        VMLA.F32 q10, q4, d1[0]
372        VMLA.F32 q11, q5, d1[0]
373        VLDM    r9!, {d12-d15}          // B1
374        VMLA.F32 q12, q4, d2[0]
375        VMLA.F32 q13, q5, d2[0]
376        VMLA.F32 q14, q4, d3[0]
377        VMLA.F32 q15, q5, d3[0]
378        VMLA.F32 q8, q6, d0[1]
379        VMLA.F32 q9, q7, d0[1]
380        VMLA.F32 q10, q6, d1[1]
381        VMLA.F32 q11, q7, d1[1]
382        VMLA.F32 q12, q6, d2[1]
383        VMLA.F32 q13, q7, d2[1]
384        VMLA.F32 q14, q6, d3[1]
385        VMLA.F32 q15, q7, d3[1]
386
387        # Is there a remainder?- 1 float of A (4 bytes)
388        TST     r5, 4
389        BEQ     4b
390
3916:
392        # Remainder- 1 float of A (4 bytes)
393        VLDM    r3!, {s0}               // A0
394        VLDM    r9!, {d8-d11}           // B0
395        VLDM    r12!, {s2}              // A1
396        VLDM    r10!, {s4}              // A2
397        VLDM    r7!, {s6}               // A3
398        VMLA.F32 q8, q4, d0[0]
399        VMLA.F32 q9, q5, d0[0]
400        VMLA.F32 q10, q4, d1[0]
401        VMLA.F32 q11, q5, d1[0]
402        VMLA.F32 q12, q4, d2[0]
403        VMLA.F32 q13, q5, d2[0]
404        VMLA.F32 q14, q4, d3[0]
405        VMLA.F32 q15, q5, d3[0]
406        B       4b
407
408        # Store odd width
4097:
410        TST     r1, 4
411        BEQ     8f
412        VST1.32 {d28-d29},  [r6]!
413        VST1.32 {d24-d25},  [r8]!
414        VMOV    q14, q15
415        VMOV    q12, q13
416        VST1.32 {d20-d21},  [r4]!
417        VST1.32 {d16-d17}, [r11]!
418        VMOV    q10, q11
419        VMOV    q8,  q9
420
4218:
422        TST     r1, 2
423        BEQ     9f
424        VST1.32 {d28},  [r6]!
425        VST1.32 {d24},  [r8]!
426        VMOV    d28, d29
427        VMOV    d24, d25
428        VST1.32 {d20},  [r4]!
429        VST1.32 {d16}, [r11]!
430        VMOV    d20, d21
431        VMOV    d16, d17
432
4339:
434        TST     r1, 1
435        BEQ     10f
436        VST1.32 {d28[0]},  [r6]!
437        VST1.32 {d24[0]},  [r8]!
438        VST1.32 {d20[0]},  [r4]!
439        VST1.32 {d16[0]}, [r11]!
440
44110:
442        VPOP    {d8-d15}
443        ADD     sp, sp, 4               // skip r3
444        POP     {r4, r5, r6, r7, r8, r9, r10, r11, pc}
445
446END_FUNCTION xnn_f32_igemm_minmax_ukernel_4x8__aarch32_neon_cortex_a55
447# LINT.ThenChange(4x8-aarch32-neon-cortex-a55.cc)
448
449#ifdef __ELF__
450.section ".note.GNU-stack","",%progbits
451#endif
452