xref: /aosp_15_r20/external/XNNPACK/src/f32-gemm/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_gemm_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
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 -> r7
21//     const union xnn_f32_minmax_params params)  sp + 116 -> (r7)
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// A0   r3  d0
27// A1  r12  d1
28// A2  r10  d2
29// A3   r0  d3
30
31// B    r9  d8,  d9, d10, d11
32// B       d12, d13, d14, d15
33
34// C0  r11 d16-d17  q8  d18-d19  q9
35// C1   r4 d20-d21 q10  d22-d23 q11
36// C2   r8 d24-d25 q12  d26-d27 q13
37// C3   r6 d28-d29 q14  d30-d31 q15
38
39// Clamp (r5) d4 d5 d6 d7
40
41BEGIN_FUNCTION xnn_f32_gemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_ld64
42        .arm
43#ifndef __APPLE__
44        .arch   armv7-a
45        .fpu    neon
46#endif
47        # Push 96 bytes
48        PUSH    {r4, r5, r6, r7, r8, r9, r10, r11}  // 32
49        VPUSH   {d8-d15}                            // +64 = 96
50
51        LDR     r7, [sp, 96]            // a_stride
52        LDR     r11, [sp, 104]          // c
53        LDR     r6, [sp, 108]           // cm_stride
54        LDR     r9, [sp, 100]           // w
55        LDR     r5, [sp, 116]           // params
56
57        # Clamp A and C pointers
58        CMP     r0, 2                   // if mr >= 2
59        ADD     r12, r3, r7             //   a1 = a0 + a_stride
60        ADD     r4, r11, r6             //   c1 = c0 + cm_stride
61        MOVLO   r12, r3                 // a1
62        MOVLO   r4, r11                 // c1
63                                        // if mr > 2
64        ADD     r10, r12, r7            //   a2 = a1 + a_stride
65        ADD     r8, r4, r6              //   c2 = c1 + cm_stride
66        MOVLS   r10, r12                // a2
67        MOVLS   r8, r4                  // c2
68
69        CMP     r0, 4                   // if mr >=4
70        ADD     r0, r10, r7             //   a3 = a2 + a_stride
71        ADD     r6, r8, r6              //   c3 = c2 + cm_stride
72        MOVLO   r0, r10                 // a3
73        MOVLO   r6, r8                  // c3
74
75        # Load min/max values
76        VLD1.32 {d4[], d5[]}, [r5]!
77        LDR     r7, [sp, 112]           // cn_stride
78        VLD1.32 {d6[], d7[]}, [r5]
79
800:
81        # Load initial bias from w into accumulators
82        VLDM    r9!, {d16-d19}          // Bias
83        SUBS    r5, r2, 8
84        VMOV    q10, q8
85        VMOV    q11, q9
86        VMOV    q12, q8
87        VMOV    q13, q9
88        VMOV    q14, q8
89        VMOV    q15, q9
90
91        $if PREFETCH:
92          PLD     [r3,  0]                // Prefetch A
93          PLD     [r3, 64]
94          PLD     [r12,  0]
95          PLD     [r12, 64]
96          PLD     [r10,  0]
97          PLD     [r10, 64]
98          PLD     [r0,  0]
99          PLD     [r0, 64]
100          PLD     [r9,   0]               // Prefetch B
101          PLD     [r9,  64]
102          PLD     [r9, 128]
103          PLD     [r9, 192]
104          PLD     [r9, 256]
105          PLD     [r9, 320]
106          PLD     [r9, 384]
107          PLD     [r9, 448]
108        BLO     3f                      // less than 2 channels?
109
110        # Main loop - 2 floats of A (8 bytes)
1111:
112        VLD1.32 {d0}, [r3]!             // A0
113        VLDM    r9!, {d8-d11}           // B0
114        VLD1.32 {d1}, [r12]!            // A1
115        VLD1.32 {d2}, [r10]!            // A2
116        VLD1.32 {d3}, [ r0]!            // A3
117        VLDM    r9!, {d12-d15}          // B1
118
119        VMLA.F32 q8, q4, d0[0]
120        VMLA.F32 q9, q5, d0[0]
121        $if PREFETCH:
122          PLD     [r3, 128]               // Prefetch A0
123        VMLA.F32 q10, q4, d1[0]
124        VMLA.F32 q13, q5, d2[0]
125        $if PREFETCH:
126          PLD     [r12, 128]              // Prefetch A1
127        VMLA.F32 q11, q5, d1[0]
128        VMLA.F32 q12, q4, d2[0]
129        $if PREFETCH:
130          PLD     [r10, 128]              // Prefetch A2
131        VMLA.F32 q14, q4, d3[0]
132        VMLA.F32 q15, q5, d3[0]
133        $if PREFETCH:
134          PLD     [r0, 128]               // Prefetch A3
135        VMLA.F32 q8, q6, d0[1]
136        VMLA.F32 q9, q7, d0[1]
137        $if PREFETCH:
138          PLD     [r9, 448]               // Prefetch B
139        VMLA.F32 q10, q6, d1[1]
140        VMLA.F32 q11, q7, d1[1]
141        SUBS    r5, r5, 8
142        VMLA.F32 q12, q6, d2[1]
143        VMLA.F32 q13, q7, d2[1]
144        VMLA.F32 q14, q6, d3[1]
145        VMLA.F32 q15, q7, d3[1]
146        BHS     1b
147
148        # Is there a remainder?- 1 float of A (4 bytes)
149        TST     r5, 4
150        BNE     3f
151
1522:
153        # Clamp
154        VMAX.F32 q8,  q8, q2
155        SUBS    r1, r1, 8
156        VMAX.F32 q9,  q9, q2
157        VMAX.F32 q10, q10, q2
158        VMAX.F32 q11, q11, q2
159        VMAX.F32 q12, q12, q2
160        VMAX.F32 q13, q13, q2
161        VMAX.F32 q14, q14, q2
162        VMAX.F32 q15, q15, q2
163        VMIN.F32 q8,  q8, q3
164        VMIN.F32 q9,  q9, q3
165        VMIN.F32 q10, q10, q3
166        VMIN.F32 q11, q11, q3
167        VMIN.F32 q12, q12, q3
168        VMIN.F32 q13, q13, q3
169        VMIN.F32 q14, q14, q3
170        VMIN.F32 q15, q15, q3
171
172        # Store full 4 x 8
173        BLO     4f
174        VST1.32 {d16-d19}, [r11], r7
175        SUB     r0, r0, r2
176        VST1.32 {d20-d23}, [r4], r7
177        SUB     r10, r10, r2
178        VST1.32 {d24-d27}, [r8], r7
179        SUB     r12, r12, r2
180        VST1.32 {d28-d31}, [r6], r7
181        SUB     r3, r3, r2
182        BHI     0b
183
184        VPOP    {d8-d15}
185        POP     {r4, r5, r6, r7, r8, r9, r10, r11}
186        BX      lr
187
1883:
189        # Remainder- 1 float of A (4 bytes)
190        VLDM    r3!,  {s0}              // A0
191        VLDM    r9!, {d8-d11}           // B0
192        VLDM    r12!, {s2}              // A1
193        VLDM    r10!, {s4}              // A2
194        VLDM    r0!, {s6}               // A3
195        VMLA.F32 q8, q4, d0[0]
196        VMLA.F32 q9, q5, d0[0]
197        VMLA.F32 q10, q4, d1[0]
198        VMLA.F32 q11, q5, d1[0]
199        VMLA.F32 q12, q4, d2[0]
200        VMLA.F32 q13, q5, d2[0]
201        VMLA.F32 q14, q4, d3[0]
202        VMLA.F32 q15, q5, d3[0]
203        B       2b
204
205        # Store odd width
2064:
207        TST     r1, 4
208        BEQ     5f
209        VST1.32 {d16-d17}, [r11]!
210        VST1.32 {d20-d21},  [r4]!
211        VMOV    q8,  q9
212        VMOV    q10, q11
213        VST1.32 {d24-d25},  [r8]!
214        VST1.32 {d28-d29},  [r6]!
215        VMOV    q12, q13
216        VMOV    q14, q15
217
2185:
219        TST     r1, 2
220        BEQ     6f
221        VST1.32 {d16}, [r11]!
222        VST1.32 {d20},  [r4]!
223        VMOV    d16, d17
224        VMOV    d20, d21
225        VST1.32 {d24},  [r8]!
226        VST1.32 {d28},  [r6]!
227        VMOV    d24, d25
228        VMOV    d28, d29
229
2306:
231        TST     r1, 1
232        BEQ     7f
233        VST1.32 {d16[0]}, [r11]
234        VST1.32 {d20[0]},  [r4]
235        VST1.32 {d24[0]},  [r8]
236        VST1.32 {d28[0]},  [r6]
237
2387:
239        VPOP    {d8-d15}
240        POP     {r4, r5, r6, r7, r8, r9, r10, r11}
241        BX      lr
242
243END_FUNCTION xnn_f32_gemm_minmax_ukernel_4x8__aarch32_neon${"_prfm" if PREFETCH else ""}_ld64
244# LINT.ThenChange(4x8-aarch32-neon-ld64.cc)
245
246#ifdef __ELF__
247.section ".note.GNU-stack","",%progbits
248#endif
249
250