xref: /aosp_15_r20/external/XNNPACK/src/qs8-gemm/4x8c4-aarch32-neondot-cortex-a55.S.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36) 
1// Copyright 2022 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$assert REQUANTIZATION in ["FP32", "RNDNU"]
7$assert not CHANNELWISE or REQUANTIZATION == "FP32"
8$assert DATATYPE in ["QC8", "QS8"]
9$assert DATATYPE != "QC8" or REQUANTIZATION == "FP32"
10
11#include <xnnpack/assembly.h>
12
13.syntax unified
14
15$PARAMS_UNION = "xnn_qs8_minmax_params" if CHANNELWISE else "xnn_qs8_conv_minmax_params"
16// void xnn_${DATATYPE.lower()}_gemm_minmax_${REQUANTIZATION.lower()}_ukernel_4x8c4__aarch32_neondot_cortex_a55(
17//     size_t mr,                            r0
18//     size_t nc,                            r1
19//     size_t kc,                            r2 -> r5
20//     const uint8_t*restrict a,             r3
21//     size_t a_stride,           sp + 80 -> (r7)
22//     const void*restrict w,     sp + 84 -> r9
23//     uint8_t*restrict c,        sp + 88 -> r11
24//     size_t cm_stride,          sp + 92 -> (r6)
25//     size_t cn_stride,          sp + 96 -> r7
26//     ${PARAMS_UNION} params)  sp + 100 -> (r5)
27
28// d8-d15, r4-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.
29
30// Register usage
31// A0   r3  d0
32// A1  r12  d1
33// A2  r10  d2
34// A3   r0  d3
35
36// B    r9  q2 q3 q4 q5
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// unused q7
44
45$if REQUANTIZATION == "RNDNU":
46  // params structure is 16 bytes
47  //  struct {
48  //    int32_t right_pre_shift;    d12[0]
49  //    int32_t multiplier;         d12[1]
50  //    int32_t right_post_shift;   d13[0]
51  //    int16_t output_zero_point;  d13[2]
52  //    int8_t output_min;          d13[6]
53  //    int8_t output_max;          d13[7]
54  //  } rndnu_neon;
55$else:
56  // params structure is 4 bytes
57  //  struct {
58  //    int16_t output_zero_point;  d13[2]
59  //    int8_t output_min;          d13[6]
60  //    int8_t output_max;          d13[7]
61  //  } xnn_qs8_minmax_params.neonv8;
62
63// iOS does not support 32 bit ARM with Neon DotProduct.
64#ifndef __APPLE__
65BEGIN_FUNCTION xnn_${DATATYPE.lower()}_gemm_minmax_${REQUANTIZATION.lower()}_ukernel_4x8c4__aarch32_neondot_cortex_a55
66        # Push 80 bytes
67        PUSH    {r4, r5, r6, r7, r8, r9, r10, r11}  // 32
68        VPUSH   {d8-d13}                            // +48 = 80
69
70        LDR     r7, [sp, 80]            // a_stride
71        ADD     r2, r2, 3               // kc = (kc + 3) & ~3
72        LDR     r11, [sp, 88]           // c
73        LDR     r6, [sp, 92]            // cm_stride
74        LDR     r9, [sp, 84]            // w
75        BIC     r2, r2, 3
76        LDR     r5, [sp, 100]           // params
77
78        # Clamp A and C pointers
79        CMP     r0, 2                   // if mr >= 2
80        ADD     r12, r3, r7             //   a1 = a0 + a_stride
81        ADD     r4, r11, r6             //   c1 = c0 + cm_stride
82        MOVLO   r12, r3                 // a1
83        MOVLO   r4, r11                 // c1
84                                        // if mr > 2
85        ADD     r10, r12, r7            //   a2 = a1 + a_stride
86        ADD     r8, r4, r6              //   c2 = c1 + cm_stride
87        MOVLS   r10, r12                // a2
88        MOVLS   r8, r4                  // c2
89
90        CMP     r0, 4                   // if mr >=4
91        ADD     r0, r10, r7             //   a3 = a2 + a_stride
92        ADD     r6, r8, r6              //   c3 = c2 + cm_stride
93        MOVLO   r0, r10                 // a3
94        MOVLO   r6, r8                  // c3
95
96        # Load params values
97        $if REQUANTIZATION == "RNDNU":
98          VLDM    r5, {d12-d13}           // RNDNU params
99        $else:
100          VLD1.32 {d13[]}, [r5]           // QC8 params
101        LDR     r7, [sp, 96]            // cn_stride
102
103        .p2align 3
1040:
105        # Load initial bias from w into accumulators
106        VLDM    r9!, {d16-d19}          // Bias
107        SUBS    r5, r2, 8               // k = kc - 8
108
109        # Prologue + Bias
110        VLD1.8  {d4},  [r9]!            // B0
111        VMOV    q10, q8
112        VLD1.8  {d0},  [r3]!            // A0
113        VMOV    q11, q9
114        VLD1.8  {d5},  [r9]!            // B1
115        VMOV    q12, q8
116        VLD1.8  {d6},  [r9]!            // B2
117        VMOV    q13, q9
118        VLD1.8  {d1}, [r12]!            // A1
119        VMOV    q14, q8
120        VLD1.8  {d7},  [r9]!            // B3
121        VMOV    q15, q9
122        BLO     4f                      // less than 8 channels?
123
124        SUBS    r5, r5, 8               // k = k - 8
125        BLO     12f                     // less than 16 channels - skip mainloop
126
127        # Main loop - 8 bytes of A.
128        # 16 SDOT, 12 LD64
129        .p2align 3
1301:
131        VSDOT.S8 q8, q2, d0[0]
132        VLD1.8  {d2}, [r10]!            // A2
133        VSDOT.S8 q9, q3, d0[0]
134        VLD1.8  {d3},  [r0]!            // A3
135        VSDOT.S8 q10, q2, d1[0]
136        VLD1.8  {d8},  [r9]!            // B4
137        VSDOT.S8 q11, q3, d1[0]
138        VLD1.8  {d9},  [r9]!            // B5
139        VSDOT.S8 q12, q2, d2[0]
140        VLD1.8  {d10},  [r9]!           // B6
141        VSDOT.S8 q13, q3, d2[0]
142        VLD1.8  {d11},  [r9]!           // B7
143        VSDOT.S8 q14, q2, d3[0]
144        VSDOT.S8 q15, q3, d3[0]
145        SUBS    r5, r5, 8
146
147        VSDOT.S8 q8, q4, d0[1]
148        VLD1.8  {d4},  [r9]!            // B0
149        VSDOT.S8 q9, q5, d0[1]
150        VLD1.8  {d5},  [r9]!            // B1
151        VSDOT.S8 q10, q4, d1[1]
152        VLD1.8  {d6},  [r9]!            // B2
153        VSDOT.S8 q11, q5, d1[1]
154        VLD1.8  {d7},  [r9]!            // B3
155        VSDOT.S8 q12, q4, d2[1]
156        VLD1.8  {d0},  [r3]!            // A0
157        VSDOT.S8 q13, q5, d2[1]
158        VLD1.8  {d1}, [r12]!            // A1
159        VSDOT.S8 q14, q4, d3[1]
160        VSDOT.S8 q15, q5, d3[1]
161        BHS     1b
162
163        # Epilogue
164        .p2align 3
16512:
166        VSDOT.S8 q8, q2, d0[0]
167        VLD1.8  {d2}, [r10]!            // A2
168        VSDOT.S8 q9, q3, d0[0]
169        VLD1.8  {d3},  [r0]!            // A3
170        VSDOT.S8 q10, q2, d1[0]
171        VLD1.8  {d8},  [r9]!            // B4
172        VSDOT.S8 q11, q3, d1[0]
173        VLD1.8  {d9},  [r9]!            // B5
174        VSDOT.S8 q12, q2, d2[0]
175        VLD1.8  {d10},  [r9]!           // B6
176        VSDOT.S8 q13, q3, d2[0]
177        VLD1.8  {d11},  [r9]!           // B7
178        VSDOT.S8 q14, q2, d3[0]
179        VSDOT.S8 q15, q3, d3[0]
180        TST     r5, 7
181
182        VSDOT.S8 q8, q4, d0[1]
183        VSDOT.S8 q9, q5, d0[1]
184        VSDOT.S8 q10, q4, d1[1]
185        VSDOT.S8 q11, q5, d1[1]
186        VSDOT.S8 q12, q4, d2[1]
187        VSDOT.S8 q13, q5, d2[1]
188        VSDOT.S8 q14, q4, d3[1]
189        VSDOT.S8 q15, q5, d3[1]
190
191        # Is there a remainder?- 4 bytes of A
192        BNE     3f
193
1942:
195        $if REQUANTIZATION == "RNDNU":
196          # RNDNU quantization
197          VDUP.32 q0, d12[0]              // right_pre_shift
198
199          VQSHL.S32 q8,  q8, q0
200          VQSHL.S32 q9,  q9, q0
201          VQSHL.S32 q10, q10, q0
202          VQSHL.S32 q11, q11, q0
203          VQSHL.S32 q12, q12, q0
204          VQSHL.S32 q13, q13, q0
205          VQSHL.S32 q14, q14, q0
206          VQSHL.S32 q15, q15, q0
207
208          VDUP.32 q2, d13[0]              // right_post_shift
209
210          VQDMULH.S32 q8,  q8, d12[1]     // multiplier
211          VQDMULH.S32 q9,  q9, d12[1]
212          VQDMULH.S32 q10, q10, d12[1]
213          VQDMULH.S32 q11, q11, d12[1]
214          VQDMULH.S32 q12, q12, d12[1]
215          VQDMULH.S32 q13, q13, d12[1]
216          VQDMULH.S32 q14, q14, d12[1]
217          VQDMULH.S32 q15, q15, d12[1]
218
219          VRSHL.S32 q8,  q8, q2
220          VRSHL.S32 q9,  q9, q2
221          VRSHL.S32 q10, q10, q2
222          VRSHL.S32 q11, q11, q2
223          VRSHL.S32 q12, q12, q2
224          VRSHL.S32 q13, q13, q2
225          VRSHL.S32 q14, q14, q2
226          VRSHL.S32 q15, q15, q2
227        $else:
228          # QC8 FP32 quantization
229          VLD1.8  {q0-q1},  [r9]!
230
231          VCVT.F32.S32 q8,  q8
232          VCVT.F32.S32 q9,  q9
233          VCVT.F32.S32 q10, q10
234          VCVT.F32.S32 q11, q11
235          VCVT.F32.S32 q12, q12
236          VCVT.F32.S32 q13, q13
237          VCVT.F32.S32 q14, q14
238          VCVT.F32.S32 q15, q15
239
240          VMUL.F32 q8,  q8, q0            // multiplier
241          VMUL.F32 q9,  q9, q1
242          VMUL.F32 q10, q10, q0
243          VMUL.F32 q11, q11, q1
244          VMUL.F32 q12, q12, q0
245          VMUL.F32 q13, q13, q1
246          VMUL.F32 q14, q14, q0
247          VMUL.F32 q15, q15, q1
248
249          VCVTN.S32.F32 q8,  q8
250          VCVTN.S32.F32 q9,  q9
251          VCVTN.S32.F32 q10, q10
252          VCVTN.S32.F32 q11, q11
253          VCVTN.S32.F32 q12, q12
254          VCVTN.S32.F32 q13, q13
255          VCVTN.S32.F32 q14, q14
256          VCVTN.S32.F32 q15, q15
257
258        VDUP.16 q0, d13[2]              // output_zero_point
259
260        VQMOVN.S32 d16, q8
261        VQMOVN.S32 d17, q9
262        VQMOVN.S32 d18, q10
263        VQMOVN.S32 d19, q11
264        VQMOVN.S32 d20, q12
265        VQMOVN.S32 d21, q13
266        VQMOVN.S32 d22, q14
267        VQMOVN.S32 d23, q15
268
269        VQADD.S16 q8,  q8, q0
270        VQADD.S16 q9,  q9, q0
271        VQADD.S16 q10, q10, q0
272        VQADD.S16 q11, q11, q0
273
274        VDUP.8  q12, d13[6]             // output_min
275
276        VQMOVN.S16 d0,  q8
277        VQMOVN.S16 d1,  q9
278        VQMOVN.S16 d2, q10
279        VQMOVN.S16 d3, q11
280
281        VDUP.8  q13, d13[7]             // output_max
282
283        VMAX.S8 q0, q0, q12
284        VMAX.S8 q1, q1, q12
285
286        SUBS    r1, r1, 8
287
288        VMIN.S8 q0, q0, q13
289        VMIN.S8 q1, q1, q13
290
291        # Store full 4 x 8
292        BLO     5f
293        VST1.8  {d0}, [r11], r7
294        SUB     r3, r3, r2
295        VST1.8  {d1}, [r4], r7
296        SUB     r12, r12, r2
297        VST1.8  {d2}, [r8], r7
298        SUB     r10, r10, r2
299        VST1.8  {d3}, [r6], r7
300        SUB     r0, r0, r2
301        BHI     0b
302
303        VPOP    {d8-d13}
304        POP     {r4, r5, r6, r7, r8, r9, r10, r11}
305        BX      lr
306
307        # Remainder prologue
308        .p2align 3
3093:
310        VLD1.8  {d4},  [r9]!            // B0
311        VLD1.8  {d0},  [r3]!            // A0
312        VLD1.8  {d5},  [r9]!            // B1
313        VLD1.8  {d6},  [r9]!            // B2
314        VLD1.8  {d1}, [r12]!            // A1
315        VLD1.8  {d7},  [r9]!            // B3
316
317        # Remainder- 4 bytes of A
3184:
319        VSDOT.S8 q8, q2, d0[0]
320        VLD1.32 {d2[0]}, [r10]!         // A2
321        VSDOT.S8 q9, q3, d0[0]
322        VLD1.32 {d3[0]},  [r0]!         // A3
323        VSDOT.S8 q10, q2, d1[0]
324        SUB     r3, r3, 4               // Rewind A0
325        VSDOT.S8 q11, q3, d1[0]
326        SUB     r12, r12, 4             // Rewind A1
327        VSDOT.S8 q12, q2, d2[0]
328        VSDOT.S8 q13, q3, d2[0]
329        VSDOT.S8 q14, q2, d3[0]
330        VSDOT.S8 q15, q3, d3[0]
331        B       2b
332
333        # Store odd width
334        .p2align 3
3355:
336        TST     r1, 4
337        BEQ     6f
338        VST1.32 {d0[0]}, [r11]!
339        VST1.32 {d1[0]}, [r4]!
340        VST1.32 {d2[0]}, [r8]!
341        VST1.32 {d3[0]}, [r6]!
342        VEXT.8  q0, q0, q0, 4
343        VEXT.8  q1, q1, q1, 4
3446:
345        TST     r1, 2
346        BEQ     7f
347        VST1.16 {d0[0]}, [r11]!
348        VST1.16 {d1[0]}, [r4]!
349        VST1.16 {d2[0]}, [r8]!
350        VST1.16 {d3[0]}, [r6]!
351        VEXT.8  q0, q0, q0, 2
352        VEXT.8  q1, q1, q1, 2
353
3547:
355        TST     r1, 1
356        BEQ     12f
357        VST1.8  {d0[0]}, [r11]
358        VST1.8  {d1[0]}, [r4]
359        VST1.8  {d2[0]}, [r8]
360        VST1.8  {d3[0]}, [r6]
361
36212:
363        VPOP    {d8-d13}
364        POP     {r4, r5, r6, r7, r8, r9, r10, r11}
365        BX      lr
366
367END_FUNCTION xnn_${DATATYPE.lower()}_gemm_minmax_${REQUANTIZATION.lower()}_ukernel_4x8c4__aarch32_neondot_cortex_a55
368#endif  // __APPLE__
369
370#ifdef __ELF__
371.section ".note.GNU-stack","",%progbits
372#endif
373
374