xref: /aosp_15_r20/external/XNNPACK/src/f32-igemm/4x8-minmax-aarch32-neon-ld64.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_igemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_ld64(
12//     size_t mr,                            r0
13//     size_t nc,                            r1
14//     size_t kc,                            r2 -> r5 -> sp + 68
15//     size_t ks,                            r3 -> sp + 72 -> r14
16//     const float**restrict a,  sp + 112 -> r2
17//     const void*restrict w,    sp + 116 -> r9
18//     uint8_t*restrict c,       sp + 120 -> r11
19//     size_t cm_stride,         sp + 124 -> (r6)
20//     size_t cn_stride,         sp + 128 -> (r7)
21//     size_t a_offset,          sp + 132 -> (r5)
22//     const float* zero,        sp + 136 -> (r7)
23//     minmax_params*params,     sp + 140 -> (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// A0   r3  d0
29// A1  r12  d1
30// A2  r10  d2
31// A3   r0  d3
32
33// B    r9  d8,  d9, d10, d11
34// B       d12, d13, d14, d15
35
36// C0  r11 d16-d17  q8  d18-d19  q9
37// C1   r4 d20-d21 q10  d22-d23 q11
38// C2   r8 d24-d25 q12  d26-d27 q13
39// C3   r6 d28-d29 q14  d30-d31 q15
40
41// Clamp (r5) d4 d5 d6 d7
42
43BEGIN_FUNCTION xnn_f32_igemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_ld64
44        .arm
45#ifndef __APPLE__
46        .arch   armv7-a
47        .fpu    neon
48#endif
49        # Push 112 bytes
50        # r2 will be reloaded in outer loop.  r3 is ks
51        PUSH    {r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, lr}   // +44
52        SUB     sp, sp, 4                                        // 4
53        VPUSH   {d8-d15}                                         // +64 = 112
54
55        LDR     r11, [sp, 120]          // c
56        LDR     r6, [sp, 124]           // cm_stride
57        LDR     r2, [sp, 112]           // a
58        LDR     r9, [sp, 116]           // w
59        LDR     r5, [sp, 140]           // params
60        MOV     r14, r3                 // p = ks
61
62        # Clamp C pointers
63        CMP     r0, 2                   // if mr >= 2
64        ADD     r4, r11, r6             //   c1 = c0 + cm_stride
65        MOVLO   r4, r11                 // c1
66                                        // if mr > 2
67        ADD     r8, r4, r6              //   c2 = c1 + cm_stride
68        MOVLS   r8, r4                  // c2
69        CMP     r0, 4                   // if mr >=4
70        ADD     r6, r8, r6              //   c3 = c2 + cm_stride
71        MOVLO   r6, r8                  // c3
72
73        # Load min/max values
74        VLD1.32 {d4[], d5[]}, [r5]!
75        VLD1.32 {d6[], d7[]}, [r5]
76
770:
78        # Load initial bias from w into accumulators
79        VLDM    r9!, {d16-d19}          // Bias
80        VMOV    q10, q8
81        VMOV    q11, q9
82        VMOV    q12, q8
83        VMOV    q13, q9
84        VMOV    q14, q8
85        VMOV    q15, q9
86
87        $if PREFETCH:
88          PLD     [r9,   0]               // Prefetch B
89          PLD     [r9,  64]
90          PLD     [r9, 128]
91          PLD     [r9, 192]
92          PLD     [r9, 256]
93          PLD     [r9, 320]
94          PLD     [r9, 384]
95          PLD     [r9, 448]
961:
97        # Load next 4 A pointers
98        LDR     r3, [r2,  0]
99        LDR     r12, [r2,  4]
100        LDR     r10, [r2,  8]
101        LDR     r0, [r2, 12]
102        ADD     r2, r2, 16
103
104        # Add a_offset
105        LDR     r5, [sp, 132]           // a_offset
106        LDR     r7, [sp, 136]           // zero
107        CMP     r3,  r7                 // if a0 == zero
108        ADD     r3,  r3, r5             // a0 += a_offset
109        MOVEQ   r3,  r7                 //   a0 = zero, else += a0 + a_offset
110        CMP     r12,  r7                // if a1 == zero
111        ADD     r12, r12, r5            // a1 += a_offset
112        MOVEQ   r12,  r7                //   a1 = zero, else += a1 + a_offset
113        CMP     r10,  r7                // if a2 == zero
114        ADD     r10, r10, r5            // a2 += a_offset
115        MOVEQ   r10,  r7                //   a2 = zero, else += a2 + a_offset
116        CMP     r0,  r7                 // if a3 == zero
117        ADD     r0,  r0, r5             // a3 += a_offset
118        LDR     r5, [sp, 68]            // kc
119        MOVEQ   r0,  r7                 //   a3 = zero, else += a3 + a_offset
120
121        $if PREFETCH:
122          PLD     [r3,  0]                // Prefetch A
123          PLD     [r3, 64]
124          PLD     [r12,  0]
125          PLD     [r12, 64]
126          PLD     [r10,  0]
127          PLD     [r10, 64]
128          PLD     [r0,  0]
129          PLD     [r0, 64]
130
131        SUBS    r5, r5, 8               // kc - 8
132        BLO     4f                      // less than 2 channels?
133
134        # Main loop - 2 floats of A (8 bytes)
1352:
136        VLD1.32 {d0}, [r3]!             // A0
137        VLDM    r9!, {d8-d11}           // B0
138        VLD1.32 {d1}, [r12]!            // A1
139        VLD1.32 {d2}, [r10]!            // A2
140        VLD1.32 {d3}, [ r0]!            // A3
141        VLDM    r9!, {d12-d15}          // B1
142
143        VMLA.F32 q8, q4, d0[0]
144        VMLA.F32 q9, q5, d0[0]
145        $if PREFETCH:
146          PLD     [r3, 128]               // Prefetch A0
147        VMLA.F32 q10, q4, d1[0]
148        VMLA.F32 q11, q5, d1[0]
149        $if PREFETCH:
150          PLD     [r12, 128]              // Prefetch A1
151        VMLA.F32 q12, q4, d2[0]
152        VMLA.F32 q13, q5, d2[0]
153        $if PREFETCH:
154          PLD     [r10, 128]              // Prefetch A2
155        VMLA.F32 q14, q4, d3[0]
156        VMLA.F32 q15, q5, d3[0]
157        $if PREFETCH:
158          PLD     [r0, 128]               // Prefetch A3
159        VMLA.F32 q8, q6, d0[1]
160        VMLA.F32 q9, q7, d0[1]
161        $if PREFETCH:
162          PLD     [r9, 448]               // Prefetch B
163        VMLA.F32 q10, q6, d1[1]
164        VMLA.F32 q11, q7, d1[1]
165        SUBS    r5, r5, 8
166        VMLA.F32 q12, q6, d2[1]
167        VMLA.F32 q13, q7, d2[1]
168        VMLA.F32 q14, q6, d3[1]
169        VMLA.F32 q15, q7, d3[1]
170        BHS     2b
171
172        # Is there a remainder?- 1 float of A (4 bytes)
173        TST     r5, 4
174        BNE     4f
175
1763:
177        # ks loop
178        SUBS    r14, r14, 16            // ks -= MR * sizeof(void*)
179        BHI     1b
180
181        LDR     r7, [sp, 128]           // cn_stride
182        LDR     r14, [sp, 72]           // p = ks
183
184        # Clamp
185        VMAX.F32 q8,  q8, q2
186        SUBS    r1, r1, 8
187        VMAX.F32 q9,  q9, q2
188        VMAX.F32 q10, q10, q2
189        VMAX.F32 q11, q11, q2
190        VMAX.F32 q12, q12, q2
191        VMAX.F32 q13, q13, q2
192        VMAX.F32 q14, q14, q2
193        VMAX.F32 q15, q15, q2
194        VMIN.F32 q8,  q8, q3
195        VMIN.F32 q9,  q9, q3
196        VMIN.F32 q10, q10, q3
197        VMIN.F32 q11, q11, q3
198        VMIN.F32 q12, q12, q3
199        VMIN.F32 q13, q13, q3
200        VMIN.F32 q14, q14, q3
201        VMIN.F32 q15, q15, q3
202
203        # Store full 4 x 8
204        BLO     5f
205        VST1.32 {d28-d31},  [r6], r7
206        VST1.32 {d24-d27},  [r8], r7
207        VST1.32 {d20-d23},  [r4], r7
208        VST1.32 {d16-d19}, [r11], r7
209        SUB     r2, r2, r14             // a -= ks
210        BHI     0b
211
212        VPOP    {d8-d15}
213        ADD     sp, sp, 12              // skip pad, r2, r3
214        POP     {r4, r5, r6, r7, r8, r9, r10, r11, pc}
215
2164:
217        # Remainder- 1 float of A (4 bytes)
218        VLDM    r3!, {s0}               // A0
219        VLDM    r9!, {d8-d11}           // B0
220        VLDM    r12!, {s2}              // A1
221        VLDM    r10!, {s4}              // A2
222        VLDM    r0!, {s6}               // A3
223        VMLA.F32 q8, q4, d0[0]
224        VMLA.F32 q9, q5, d0[0]
225        VMLA.F32 q10, q4, d1[0]
226        VMLA.F32 q11, q5, d1[0]
227        VMLA.F32 q12, q4, d2[0]
228        VMLA.F32 q13, q5, d2[0]
229        VMLA.F32 q14, q4, d3[0]
230        VMLA.F32 q15, q5, d3[0]
231        B       3b
232
233        # Store odd width
2345:
235        TST     r1, 4
236        BEQ     6f
237        VST1.32 {d28-d29},  [r6]!
238        VST1.32 {d24-d25},  [r8]!
239        VMOV    q14, q15
240        VMOV    q12, q13
241        VST1.32 {d20-d21},  [r4]!
242        VST1.32 {d16-d17}, [r11]!
243        VMOV    q10, q11
244        VMOV    q8,  q9
245
2466:
247        TST     r1, 2
248        BEQ     7f
249        VST1.32 {d28},  [r6]!
250        VST1.32 {d24},  [r8]!
251        VMOV    d28, d29
252        VMOV    d24, d25
253        VST1.32 {d20},  [r4]!
254        VST1.32 {d16}, [r11]!
255        VMOV    d20, d21
256        VMOV    d16, d17
257
2587:
259        TST     r1, 1
260        BEQ     8f
261        VST1.32 {d28[0]},  [r6]!
262        VST1.32 {d24[0]},  [r8]!
263        VST1.32 {d20[0]},  [r4]!
264        VST1.32 {d16[0]}, [r11]!
265
2668:
267        VPOP    {d8-d15}
268        ADD     sp, sp, 12              // skip pad, r2, r3
269        POP     {r4, r5, r6, r7, r8, r9, r10, r11, pc}
270
271END_FUNCTION xnn_f32_igemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_ld64
272# LINT.ThenChange(4x8-aarch32-neon-ld64.cc)
273
274#ifdef __ELF__
275.section ".note.GNU-stack","",%progbits
276#endif
277