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