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
7 #include <xnnpack/aarch32-assembler.h>
8 #include <xnnpack/allocator.h>
9 #include <xnnpack/igemm.h>
10
11 namespace xnnpack {
12 namespace aarch32 {
13 namespace {
14 class Generator : public Assembler {
15 using Assembler::Assembler;
16 public:
17 void generate(size_t nc, size_t kc, size_t ks, const void* params);
18 };
19
20
21 // void xnn_qc8_igemm_minmax_fp32_ukernel_4x8c4__aarch32_neondot_ld64(
22 // size_t mr, r0
23 // size_t nc, r1
24 // size_t kc, r2 -> r5 -> sp + 52
25 // size_t ks, r3 -> sp + 56 -> r14
26 // const int8_t**restrict a, sp + 96 -> r2
27 // const void*restrict w, sp + 100 -> r9
28 // int8_t*restrict c, sp + 104 -> r11
29 // size_t cm_stride, sp + 108 -> (r6)
30 // size_t cn_stride, sp + 112 -> (r7)
31 // size_t a_offset, sp + 116 -> (r5)
32 // const int8_t* zero, sp + 120 -> (r7)
33 // xnn_qs8_minmax_params*params); sp + 124 -> (r5)
34
35 // d8-d15, r4-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.
36
37 // Register usage
38
39 // A0 r3 d0
40 // A1 r12 d1
41 // A2 r10 d2
42 // A3 r0 d3
43
44 // B r9 q2 q3 q4 q5
45
46 // C0 r11 d16-d17 q8 d18-d19 q9
47 // C1 r4 d20-d21 q10 d22-d23 q11
48 // C2 r8 d24-d25 q12 d26-d27 q13
49 // C3 r6 d28-d29 q14 d30-d31 q15
50
51 // unused q7
52
53 // params structure is 4 bytes
54 // struct {
55 // int16_t output_zero_point; d13[2]
56 // int8_t output_min; d13[6]
57 // int8_t output_max; d13[7]
58 // } xnn_qs8_minmax_params.neonv8;
59
generate(size_t nc,size_t kc,size_t ks,const void * params)60 void Generator::generate(size_t nc, size_t kc, size_t ks, const void* params) {
61 Label l0, l1, l2, l3, l4, l5, l6, l7, l8;
62
63 // Auto-generated file. Do not edit!
64 add(r2, r2, 3); // kc = (kc + 3) & ~3
65 bic(r2, r2, 3);
66
67 // Push 96 bytes
68 // r2 will be reloaded in outer loop. r3 is ks
69 push({r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, lr}); // +44
70 sub(sp, sp, 4); // 4
71 vpush({d8-d13}); // +48 = 96
72
73 ldr(r11, mem[sp, 104]); // c
74 ldr(r6, mem[sp, 108]); // cm_stride
75 ldr(r2, mem[sp, 96]); // a
76 ldr(r9, mem[sp, 100]); // w
77 ldr(r5, mem[sp, 124]); // params
78 mov(r14, r3); // p = ks
79
80 // Clamp C pointers
81 cmp(r0, 2); // if mr >= 2
82 add(r4, r11, r6); // c1 = c0 + cm_stride
83 movlo(r4, r11); // c1
84 // if mr > 2
85 add(r8, r4, r6); // c2 = c1 + cm_stride
86 movls(r8, r4); // c2
87 cmp(r0, 4); // if mr >=4
88 add(r6, r8, r6); // c3 = c2 + cm_stride
89 movlo(r6, r8); // c3
90
91 // Load params values
92 vld1r_32({d13}, mem[r5]); // QC8 params
93
94 bind(l0);
95 // Load initial bias from w into accumulators
96 vldm(mem[r9]++, {d16-d19}); // Bias
97 vmov(q10, q8);
98 vmov(q11, q9);
99 vmov(q12, q8);
100 vmov(q13, q9);
101 vmov(q14, q8);
102 vmov(q15, q9);
103
104 bind(l1);
105 // Load next 4 A pointers
106 ldr(r3, mem[r2, 0]);
107 ldr(r12, mem[r2, 4]);
108 ldr(r10, mem[r2, 8]);
109 ldr(r0, mem[r2, 12]);
110 add(r2, r2, 16);
111
112 // Add a_offset
113 ldr(r5, mem[sp, 116]); // a_offset
114 ldr(r7, mem[sp, 120]); // zero
115 cmp(r3, r7); // if a0 == zero
116 add(r3, r3, r5); // a0 += a_offset
117 moveq(r3, r7); // a0 = zero, else += a0 + a_offset
118 cmp(r12, r7); // if a1 == zero
119 add(r12, r12, r5); // a1 += a_offset
120 moveq(r12, r7); // a1 = zero, else += a1 + a_offset
121 cmp(r10, r7); // if a2 == zero
122 add(r10, r10, r5); // a2 += a_offset
123 moveq(r10, r7); // a2 = zero, else += a2 + a_offset
124 cmp(r0, r7); // if a3 == zero
125 add(r0, r0, r5); // a3 += a_offset
126 ldr(r5, mem[sp, 52]); // kc
127 moveq(r0, r7); // a3 = zero, else += a3 + a_offset
128
129 subs(r5, r5, 8); // kc - 8
130 blo(l4); // less than 8 channels?
131
132 // Main loop - 8 bytes of A.
133 // 16 SDOT, 4 LD64 A, 4 LD128 B
134 align(8);
135 bind(l2);
136 vld1_8({d0}, mem[r3]++); // A0
137 vld1_8({q2}, mem[r9]++); // B0
138 vld1_8({d1}, mem[r12]++); // A1
139 vld1_8({q3}, mem[r9]++); // B1
140 vld1_8({d2}, mem[r10]++); // A2
141 vld1_8({q4}, mem[r9]++); // B2
142 vld1_8({d3}, mem[r0]++); // A3
143 vld1_8({q5}, mem[r9]++); // B3
144 subs(r5, r5, 8);
145
146 vsdot_s8(q8, q2, d0[0]);
147 vsdot_s8(q9, q3, d0[0]);
148 vsdot_s8(q10, q2, d1[0]);
149 vsdot_s8(q11, q3, d1[0]);
150 vsdot_s8(q12, q2, d2[0]);
151 vsdot_s8(q13, q3, d2[0]);
152 vsdot_s8(q14, q2, d3[0]);
153 vsdot_s8(q15, q3, d3[0]);
154
155 vsdot_s8(q8, q4, d0[1]);
156 vsdot_s8(q9, q5, d0[1]);
157 vsdot_s8(q10, q4, d1[1]);
158 vsdot_s8(q11, q5, d1[1]);
159 vsdot_s8(q12, q4, d2[1]);
160 vsdot_s8(q13, q5, d2[1]);
161 vsdot_s8(q14, q4, d3[1]);
162 vsdot_s8(q15, q5, d3[1]);
163 bhs(l2);
164
165 // Is there a remainder?- 4 bytes of A
166 tst(r5, 4);
167 bne(l4);
168
169 bind(l3);
170 // ks loop
171 subs(r14, r14, 16); // ks -= MR * sizeof(void*)
172 bhi(l1);
173
174 ldr(r7, mem[sp, 112]); // cn_stride
175 ldr(r14, mem[sp, 56]); // p = ks
176
177 // QC8 FP32 quantization
178 vld1_8({q0-q1}, mem[r9]++);
179
180 vcvt_f32_s32(q8, q8);
181 vcvt_f32_s32(q9, q9);
182 vcvt_f32_s32(q10, q10);
183 vcvt_f32_s32(q11, q11);
184 vcvt_f32_s32(q12, q12);
185 vcvt_f32_s32(q13, q13);
186 vcvt_f32_s32(q14, q14);
187 vcvt_f32_s32(q15, q15);
188
189 vmul_f32(q8, q8, q0); // multiplier
190 vmul_f32(q9, q9, q1);
191 vmul_f32(q10, q10, q0);
192 vmul_f32(q11, q11, q1);
193 vmul_f32(q12, q12, q0);
194 vmul_f32(q13, q13, q1);
195 vmul_f32(q14, q14, q0);
196 vmul_f32(q15, q15, q1);
197
198 vcvtn_s32_f32(q8, q8);
199 vcvtn_s32_f32(q9, q9);
200 vcvtn_s32_f32(q10, q10);
201 vcvtn_s32_f32(q11, q11);
202 vcvtn_s32_f32(q12, q12);
203 vcvtn_s32_f32(q13, q13);
204 vcvtn_s32_f32(q14, q14);
205 vcvtn_s32_f32(q15, q15);
206 vdup_16(q0, d13[2]); // output_zero_point
207
208 vqmovn_s32(d16, q8);
209 vqmovn_s32(d17, q9);
210 vqmovn_s32(d18, q10);
211 vqmovn_s32(d19, q11);
212 vqmovn_s32(d20, q12);
213 vqmovn_s32(d21, q13);
214 vqmovn_s32(d22, q14);
215 vqmovn_s32(d23, q15);
216
217 vqadd_s16(q8, q8, q0);
218 vqadd_s16(q9, q9, q0);
219 vqadd_s16(q10, q10, q0);
220 vqadd_s16(q11, q11, q0);
221
222 vdup_8(q12, d13[6]); // output_min
223
224 vqmovn_s16(d0, q8);
225 vqmovn_s16(d1, q9);
226 vqmovn_s16(d2, q10);
227 vqmovn_s16(d3, q11);
228
229 vdup_8(q13, d13[7]); // output_min
230
231 vmax_s8(q0, q0, q12);
232 vmax_s8(q1, q1, q12);
233
234 subs(r1, r1, 8); // nc -= 8
235
236 vmin_s8(q0, q0, q13);
237 vmin_s8(q1, q1, q13);
238
239 // Store full 4 x 8
240 blo(l5);
241 vst1_8({d3}, mem[r6], r7);
242 vst1_8({d2}, mem[r8], r7);
243 vst1_8({d1}, mem[r4], r7);
244 vst1_8({d0}, mem[r11], r7);
245 sub(r2, r2, r14); // a -= ks
246 bhi(l0);
247
248 vpop({d8-d13});
249 add(sp, sp, 12); // skip pad, r2, r3
250 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
251
252 bind(l4);
253 // Remainder- 4 bytes of A
254 vld1_32({d0[0]}, mem[r3]++); // A0
255 vld1_32({q2}, mem[r9]++); // B0
256 vld1_32({d1[0]}, mem[r12]++); // A1
257 vld1_32({q3}, mem[r9]++); // B1
258 vld1_32({d2[0]}, mem[r10]++); // A2
259 vld1_32({d3[0]}, mem[r0]++); // A3
260
261 vsdot_s8(q8, q2, d0[0]);
262 vsdot_s8(q9, q3, d0[0]);
263 vsdot_s8(q10, q2, d1[0]);
264 vsdot_s8(q11, q3, d1[0]);
265 vsdot_s8(q12, q2, d2[0]);
266 vsdot_s8(q13, q3, d2[0]);
267 vsdot_s8(q14, q2, d3[0]);
268 vsdot_s8(q15, q3, d3[0]);
269 b(l3);
270
271 // Store odd width
272 align(8);
273 bind(l5);
274 tst(r1, 4);
275 beq(l6);
276 vst1_32({d3[0]}, mem[r6]++);
277 vst1_32({d2[0]}, mem[r8]++);
278 vst1_32({d1[0]}, mem[r4]++);
279 vst1_32({d0[0]}, mem[r11]++);
280 vext_8(q0, q0, q0, 4);
281 vext_8(q1, q1, q1, 4);
282 bind(l6);
283 tst(r1, 2);
284 beq(l7);
285 vst1_16({d3[0]}, mem[r6]++);
286 vst1_16({d2[0]}, mem[r8]++);
287 vst1_16({d1[0]}, mem[r4]++);
288 vst1_16({d0[0]}, mem[r11]++);
289 vext_8(q0, q0, q0, 2);
290 vext_8(q1, q1, q1, 2);
291
292 bind(l7);
293 tst(r1, 1);
294 beq(l8);
295 vst1_8({d3[0]}, mem[r6]);
296 vst1_8({d2[0]}, mem[r8]);
297 vst1_8({d1[0]}, mem[r4]);
298 vst1_8({d0[0]}, mem[r11]);
299
300 bind(l8);
301 vpop({d8-d13});
302 add(sp, sp, 12); // skip pad, r2, r3
303 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
304 }
305 } // namespace
306 } // aarch32
307 } // xnnpack
308
xnn_generate_qc8_igemm_fp32_ukernel_4x8c4__aarch32_neondot_ld64(xnn_code_buffer * code,size_t nc,size_t kc,size_t ks,const void * params)309 xnn_status xnn_generate_qc8_igemm_fp32_ukernel_4x8c4__aarch32_neondot_ld64(xnn_code_buffer* code, size_t nc, size_t kc, size_t ks, const void* params) {
310 using namespace xnnpack::aarch32;
311 Generator g(code);
312 g.generate(nc, kc, ks, nullptr);
313 g.finalize();
314 if (g.error() != xnnpack::Error::kNoError) {
315 return xnn_status_invalid_state;
316 }
317 return xnn_status_success;
318 }
319