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