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 <cassert>
8
9 #include <xnnpack.h>
10 #include <xnnpack/aarch32-assembler.h>
11 #include <xnnpack/allocator.h>
12 #include <xnnpack/igemm.h>
13
14 namespace xnnpack {
15 namespace aarch32 {
16 namespace {
17 class Generator : public Assembler {
18 using Assembler::Assembler;
19 public:
20 void generate(bool prefetch, size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params);
21 };
22
23
24 // void xnn_qc8_igemm_minmax_fp32_ukernel_4x8__aarch32_neonv8_mlal_lane_prfm_ld64
25 // size_t mr, r0
26 // size_t nc, r1
27 // size_t kc, r2 -> r5 -> sp + 36
28 // size_t ks, r3 -> sp + 40 -> r14
29 // const int8_t**restrict a, sp + 80 -> r2
30 // const void*restrict w, sp + 84 -> r9
31 // int8_t*restrict c, sp + 88 -> r11
32 // size_t cm_stride, sp + 92 -> (r6)
33 // size_t cn_stride, sp + 96 -> (r7)
34 // size_t a_offset, sp + 100 -> (r5)
35 // const int8_t* zero, sp + 104 -> (r7)
36 // xnn_qs8_minmax_params*params); sp + 108 -> (r5)
37
38 // d8-d15, r4-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.
39
40 // Register usage
41
42 // A0 r3 d0-d1 q0
43 // A1 r12 d2-d3 q1
44 // A2 r10 d4-d5 q2
45 // A3 r0 d6-d7 q3
46
47 // B r9 d8-d9 q4
48
49 // C0 r11 d16-d17 q8 d18-d19 q9
50 // C1 r4 d20-d21 q10 d22-d23 q11
51 // C2 r8 d24-d25 q12 d26-d27 q13
52 // C3 r6 d28-d29 q14 d30-d31 q15
53
54 // Unused q6 q7
55
56 // params structure is 4 bytes
57 // struct {
58 // int16_t output_zero_point; d11[2]
59 // int8_t output_min; d11[6]
60 // int8_t output_max; d11[7]
61 // } xnn_qs8_minmax_params.neonv8;
62
63 // Converted from: src/qc8-igemm/gen/4x8-minmax-fp32-aarch32-neonv8-mlal-lane-prfm-ld64.S
generate(bool prefetch,size_t max_mr,size_t nc_mod_nr,size_t kc,size_t ks,const void * params)64 void Generator::generate(bool prefetch, size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params)
65 {
66 assert(nc_mod_nr < 8);
67 assert(kc != 0);
68 assert(ks != 0);
69
70 Label l0, l1, l2, l3, l4, l5, l6, l7, l8;
71
72 // Push 80 bytes
73 // r2 will be reloaded in outer loop. r3 is ks
74 push({r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, lr}); // +44
75 sub(sp, sp, 4); // 4
76 vpush({d8-d11}); // +32 = 80
77
78 ldr(r11, mem[sp, 88]); // c
79 ldr(r6, mem[sp, 92]); // cm_stride
80 ldr(r2, mem[sp, 80]); // a
81 ldr(r9, mem[sp, 84]); // w
82 ldr(r5, mem[sp, 108]); // params
83 mov(r14, r3); // p = ks
84
85 // Clamp C pointers
86 cmp(r0, 2); // if mr >= 2
87 add(r4, r11, r6); // c1 = c0 + cm_stride
88 movlo(r4, r11); // c1
89 // if mr > 2
90 add(r8, r4, r6); // c2 = c1 + cm_stride
91 movls(r8, r4); // c2
92 cmp(r0, 4); // if mr >=4
93 add(r6, r8, r6); // c3 = c2 + cm_stride
94 movlo(r6, r8); // c3
95
96 // Load params values
97 vld1r_32({d11}, mem[r5]); // QC8 params
98
99 if (prefetch) {
100 pld(mem[r9, 64]); // Prefetch B
101 pld(mem[r9, 128]);
102 pld(mem[r9, 192]);
103 pld(mem[r9, 256]);
104 pld(mem[r9, 320]);
105 pld(mem[r9, 384]);
106 }
107
108 align(8);
109 bind(l0);
110 // Load initial bias from w into accumulators
111 vldm(mem[r9]++, {d16-d19}); // Bias
112 vmov(q10, q8);
113 vmov(q11, q9);
114 vmov(q12, q8);
115 vmov(q13, q9);
116 vmov(q14, q8);
117 vmov(q15, q9);
118
119 align(8);
120 bind(l1);
121 // Load next 4 A pointers
122 ldr(r3, mem[r2, 0]);
123 ldr(r12, mem[r2, 4]);
124 ldr(r10, mem[r2, 8]);
125 ldr(r0, mem[r2, 12]);
126 add(r2, r2, 16);
127
128 if (prefetch) {
129 pld(mem[r3, 64]);
130 pld(mem[r12, 64]);
131 pld(mem[r10, 64]);
132 pld(mem[r0, 64]);
133 }
134
135 // Add a_offset
136 ldr(r5, mem[sp, 100]); // a_offset
137 ldr(r7, mem[sp, 104]); // zero
138 cmp(r3, r7); // if a0 == zero
139 add(r3, r3, r5); // a0 += a_offset
140 moveq(r3, r7); // a0 = zero, else += a0 + a_offset
141 cmp(r12, r7); // if a1 == zero
142 add(r12, r12, r5); // a1 += a_offset
143 moveq(r12, r7); // a1 = zero, else += a1 + a_offset
144 cmp(r10, r7); // if a2 == zero
145 add(r10, r10, r5); // a2 += a_offset
146 moveq(r10, r7); // a2 = zero, else += a2 + a_offset
147 cmp(r0, r7); // if a3 == zero
148 add(r0, r0, r5); // a3 += a_offset
149 ldr(r5, mem[sp, 36]); // kc
150 moveq(r0, r7); // a3 = zero, else += a3 + a_offset
151
152 subs(r5, r5, 8); // kc - 8
153 blo(l4); // less than 8 channels?
154
155 // Main loop - 8 bytes
156 // 64 bytes for weights.
157 align(8);
158 bind(l2);
159 vld1_8({d0}, mem[r3]++); // A0
160 vld1_8({d8}, mem[r9]++); // B
161 vld1_8({d2}, mem[r12]++); // A1
162 vld1_8({d4}, mem[r10]++); // A2
163 vld1_8({d6}, mem[r0]++); // A3
164 subs(r5, r5, 8);
165 if (prefetch) {
166 pld(mem[r3, 128]);
167 }
168 vmovl_s8(q0, d0);
169 if (prefetch) {
170 pld(mem[r12, 128]);
171 }
172 vmovl_s8(q4, d8);
173 if (prefetch) {
174 pld(mem[r10, 128]);
175 }
176 vmovl_s8(q1, d2);
177 if (prefetch) {
178 pld(mem[r0, 128]);
179 }
180 vmovl_s8(q2, d4);
181 if (prefetch) {
182 pld(mem[r9, 448]);
183 }
184 vmovl_s8(q3, d6);
185 vmlal_s16(q8, d8, d0[0]);
186 vmlal_s16(q9, d9, d0[0]);
187 vmlal_s16(q10, d8, d2[0]);
188 vmlal_s16(q11, d9, d2[0]);
189 vmlal_s16(q12, d8, d4[0]);
190 vmlal_s16(q13, d9, d4[0]);
191 vmlal_s16(q14, d8, d6[0]);
192 vmlal_s16(q15, d9, d6[0]);
193
194 vld1_8({d8}, mem[r9]++);
195 vmovl_s8(q4, d8);
196 vmlal_s16(q8, d8, d0[1]);
197 vmlal_s16(q9, d9, d0[1]);
198 vmlal_s16(q10, d8, d2[1]);
199 vmlal_s16(q11, d9, d2[1]);
200 vmlal_s16(q12, d8, d4[1]);
201 vmlal_s16(q13, d9, d4[1]);
202 vmlal_s16(q14, d8, d6[1]);
203 vmlal_s16(q15, d9, d6[1]);
204
205 vld1_8({d8}, mem[r9]++);
206 vmovl_s8(q4, d8);
207 vmlal_s16(q8, d8, d0[2]);
208 vmlal_s16(q9, d9, d0[2]);
209 vmlal_s16(q10, d8, d2[2]);
210 vmlal_s16(q11, d9, d2[2]);
211 vmlal_s16(q12, d8, d4[2]);
212 vmlal_s16(q13, d9, d4[2]);
213 vmlal_s16(q14, d8, d6[2]);
214 vmlal_s16(q15, d9, d6[2]);
215
216 vld1_8({d8}, mem[r9]++);
217 vmovl_s8(q4, d8);
218 vmlal_s16(q8, d8, d0[3]);
219 vmlal_s16(q9, d9, d0[3]);
220 vmlal_s16(q10, d8, d2[3]);
221 vmlal_s16(q11, d9, d2[3]);
222 vmlal_s16(q12, d8, d4[3]);
223 vmlal_s16(q13, d9, d4[3]);
224 vmlal_s16(q14, d8, d6[3]);
225 vmlal_s16(q15, d9, d6[3]);
226
227 vld1_8({d8}, mem[r9]++);
228 vmovl_s8(q4, d8);
229 vmlal_s16(q8, d8, d1[0]);
230 vmlal_s16(q9, d9, d1[0]);
231 vmlal_s16(q10, d8, d3[0]);
232 vmlal_s16(q11, d9, d3[0]);
233 vmlal_s16(q12, d8, d5[0]);
234 vmlal_s16(q13, d9, d5[0]);
235 vmlal_s16(q14, d8, d7[0]);
236 vmlal_s16(q15, d9, d7[0]);
237
238 vld1_8({d8}, mem[r9]++);
239 vmovl_s8(q4, d8);
240 vmlal_s16(q8, d8, d1[1]);
241 vmlal_s16(q9, d9, d1[1]);
242 vmlal_s16(q10, d8, d3[1]);
243 vmlal_s16(q11, d9, d3[1]);
244 vmlal_s16(q12, d8, d5[1]);
245 vmlal_s16(q13, d9, d5[1]);
246 vmlal_s16(q14, d8, d7[1]);
247 vmlal_s16(q15, d9, d7[1]);
248
249 vld1_8({d8}, mem[r9]++);
250 vmovl_s8(q4, d8);
251 vmlal_s16(q8, d8, d1[2]);
252 vmlal_s16(q9, d9, d1[2]);
253 vmlal_s16(q10, d8, d3[2]);
254 vmlal_s16(q11, d9, d3[2]);
255 vmlal_s16(q12, d8, d5[2]);
256 vmlal_s16(q13, d9, d5[2]);
257 vmlal_s16(q14, d8, d7[2]);
258 vmlal_s16(q15, d9, d7[2]);
259
260 vld1_8({d8}, mem[r9]++);
261 vmovl_s8(q4, d8);
262 vmlal_s16(q8, d8, d1[3]);
263 vmlal_s16(q9, d9, d1[3]);
264 vmlal_s16(q10, d8, d3[3]);
265 vmlal_s16(q11, d9, d3[3]);
266 vmlal_s16(q12, d8, d5[3]);
267 vmlal_s16(q13, d9, d5[3]);
268 vmlal_s16(q14, d8, d7[3]);
269 vmlal_s16(q15, d9, d7[3]);
270 bhs(l2);
271
272 // Is there a remainder?- 1-7 bytes of A
273 adds(r5, r5, 8);
274 bne(l4);
275
276 bind(l3);
277 // ks loop
278 subs(r14, r14, 16); // ks -= MR * sizeof(void*)
279 bhi(l1);
280
281 ldr(r7, mem[sp, 96]); // cn_stride
282 ldr(r14, mem[sp, 40]); // p = ks
283
284 // QC8 FP32 quantization
285 vld1_8({q0-q1}, mem[r9]++);
286
287 vcvt_f32_s32(q8, q8);
288 vcvt_f32_s32(q9, q9);
289 vcvt_f32_s32(q10, q10);
290 vcvt_f32_s32(q11, q11);
291 vcvt_f32_s32(q12, q12);
292 vcvt_f32_s32(q13, q13);
293 vcvt_f32_s32(q14, q14);
294 vcvt_f32_s32(q15, q15);
295
296 vmul_f32(q8, q8, q0); // multiplier
297 vmul_f32(q9, q9, q1);
298 vmul_f32(q10, q10, q0);
299 vmul_f32(q11, q11, q1);
300 vmul_f32(q12, q12, q0);
301 vmul_f32(q13, q13, q1);
302 vmul_f32(q14, q14, q0);
303 vmul_f32(q15, q15, q1);
304
305 vcvtn_s32_f32(q8, q8);
306 vcvtn_s32_f32(q9, q9);
307 vcvtn_s32_f32(q10, q10);
308 vcvtn_s32_f32(q11, q11);
309 vcvtn_s32_f32(q12, q12);
310 vcvtn_s32_f32(q13, q13);
311 vcvtn_s32_f32(q14, q14);
312 vcvtn_s32_f32(q15, q15);
313
314 vdup_16(q0, d11[2]); // output_zero_point
315
316 vqmovn_s32(d16, q8);
317 vqmovn_s32(d17, q9);
318 vqmovn_s32(d18, q10);
319 vqmovn_s32(d19, q11);
320 vqmovn_s32(d20, q12);
321 vqmovn_s32(d21, q13);
322 vqmovn_s32(d22, q14);
323 vqmovn_s32(d23, q15);
324
325 vqadd_s16(q8, q8, q0);
326 vqadd_s16(q9, q9, q0);
327 vqadd_s16(q10, q10, q0);
328 vqadd_s16(q11, q11, q0);
329
330 vdup_8(q12, d11[6]); // output_min
331
332 vqmovn_s16(d0, q8);
333 vqmovn_s16(d1, q9);
334 vqmovn_s16(d2, q10);
335 vqmovn_s16(d3, q11);
336
337 vdup_8(q13, d11[7]); // output_max
338
339 vmax_s8(q0, q0, q12);
340 vmax_s8(q1, q1, q12);
341
342 subs(r1, r1, 8); // nc -= 8
343
344 vmin_s8(q0, q0, q13);
345 vmin_s8(q1, q1, q13);
346
347 // Store full 4 x 8
348 blo(l5);
349 vst1_8({d3}, mem[r6], r7);
350 vst1_8({d2}, mem[r8], r7);
351 vst1_8({d1}, mem[r4], r7);
352 vst1_8({d0}, mem[r11], r7);
353 sub(r2, r2, r14); // a -= ks
354 bhi(l0);
355
356 vpop({d8-d11});
357 add(sp, sp, 12); // skip pad, r2, r3
358 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
359
360 // Remainder- 1 to 7 bytes of A
361 align(8);
362 bind(l4);
363 and_(r5, r5, 7); // kc remainder 1 to 7
364
365 vld1_8({d0}, mem[r3]);
366 vld1_8({d8}, mem[r9]++);
367 vld1_8({d2}, mem[r12]);
368 vld1_8({d4}, mem[r10]);
369 vld1_8({d6}, mem[r0]);
370
371 vmovl_s8(q0, d0);
372 vmovl_s8(q4, d8);
373 vmovl_s8(q1, d2);
374 vmovl_s8(q2, d4);
375 vmovl_s8(q3, d6);
376 vmlal_s16(q8, d8, d0[0]);
377 vmlal_s16(q9, d9, d0[0]);
378 vmlal_s16(q10, d8, d2[0]);
379 vmlal_s16(q11, d9, d2[0]);
380 vmlal_s16(q12, d8, d4[0]);
381 vmlal_s16(q13, d9, d4[0]);
382 vmlal_s16(q14, d8, d6[0]);
383 vmlal_s16(q15, d9, d6[0]);
384 cmp(r5, 2);
385 blo(l3);
386
387 vld1_8({d8}, mem[r9]++);
388 vmovl_s8(q4, d8);
389 vmlal_s16(q8, d8, d0[1]);
390 vmlal_s16(q9, d9, d0[1]);
391 vmlal_s16(q10, d8, d2[1]);
392 vmlal_s16(q11, d9, d2[1]);
393 vmlal_s16(q12, d8, d4[1]);
394 vmlal_s16(q13, d9, d4[1]);
395 vmlal_s16(q14, d8, d6[1]);
396 vmlal_s16(q15, d9, d6[1]);
397 beq(l3);
398
399 vld1_8({d8}, mem[r9]++);
400 vmovl_s8(q4, d8);
401 vmlal_s16(q8, d8, d0[2]);
402 vmlal_s16(q9, d9, d0[2]);
403 vmlal_s16(q10, d8, d2[2]);
404 vmlal_s16(q11, d9, d2[2]);
405 vmlal_s16(q12, d8, d4[2]);
406 vmlal_s16(q13, d9, d4[2]);
407 vmlal_s16(q14, d8, d6[2]);
408 vmlal_s16(q15, d9, d6[2]);
409 cmp(r5, 4);
410 blo(l3);
411
412 vld1_8({d8}, mem[r9]++);
413 vmovl_s8(q4, d8);
414 vmlal_s16(q8, d8, d0[3]);
415 vmlal_s16(q9, d9, d0[3]);
416 vmlal_s16(q10, d8, d2[3]);
417 vmlal_s16(q11, d9, d2[3]);
418 vmlal_s16(q12, d8, d4[3]);
419 vmlal_s16(q13, d9, d4[3]);
420 vmlal_s16(q14, d8, d6[3]);
421 vmlal_s16(q15, d9, d6[3]);
422 beq(l3);
423
424 vld1_8({d8}, mem[r9]++);
425 vmovl_s8(q4, d8);
426 vmlal_s16(q8, d8, d1[0]);
427 vmlal_s16(q9, d9, d1[0]);
428 vmlal_s16(q10, d8, d3[0]);
429 vmlal_s16(q11, d9, d3[0]);
430 vmlal_s16(q12, d8, d5[0]);
431 vmlal_s16(q13, d9, d5[0]);
432 vmlal_s16(q14, d8, d7[0]);
433 vmlal_s16(q15, d9, d7[0]);
434 cmp(r5, 6);
435 blo(l3);
436
437 vld1_8({d8}, mem[r9]++);
438 vmovl_s8(q4, d8);
439 vmlal_s16(q8, d8, d1[1]);
440 vmlal_s16(q9, d9, d1[1]);
441 vmlal_s16(q10, d8, d3[1]);
442 vmlal_s16(q11, d9, d3[1]);
443 vmlal_s16(q12, d8, d5[1]);
444 vmlal_s16(q13, d9, d5[1]);
445 vmlal_s16(q14, d8, d7[1]);
446 vmlal_s16(q15, d9, d7[1]);
447 beq(l3);
448
449 vld1_8({d8}, mem[r9]++);
450 vmovl_s8(q4, d8);
451 vmlal_s16(q8, d8, d1[2]);
452 vmlal_s16(q9, d9, d1[2]);
453 vmlal_s16(q10, d8, d3[2]);
454 vmlal_s16(q11, d9, d3[2]);
455 vmlal_s16(q12, d8, d5[2]);
456 vmlal_s16(q13, d9, d5[2]);
457 vmlal_s16(q14, d8, d7[2]);
458 vmlal_s16(q15, d9, d7[2]);
459 b(l3);
460
461 // Store odd width
462 align(8);
463 bind(l5);
464 tst(r1, 4);
465 beq(l6);
466 vst1_32({d3[0]}, mem[r6]++);
467 vst1_32({d2[0]}, mem[r8]++);
468 vst1_32({d1[0]}, mem[r4]++);
469 vst1_32({d0[0]}, mem[r11]++);
470 vext_8(q1, q1, q1, 4);
471 vext_8(q0, q0, q0, 4);
472 bind(l6);
473 tst(r1, 2);
474 beq(l7);
475 vst1_16({d3[0]}, mem[r6]++);
476 vst1_16({d2[0]}, mem[r8]++);
477 vst1_16({d1[0]}, mem[r4]++);
478 vst1_16({d0[0]}, mem[r11]++);
479 vext_8(q1, q1, q1, 2);
480 vext_8(q0, q0, q0, 2);
481
482 bind(l7);
483 tst(r1, 1);
484 beq(l8);
485 vst1_8({d3[0]}, mem[r6]);
486 vst1_8({d2[0]}, mem[r8]);
487 vst1_8({d1[0]}, mem[r4]);
488 vst1_8({d0[0]}, mem[r11]);
489
490 bind(l8);
491 vpop({d8-d11});
492 add(sp, sp, 12); // skip pad, r2, r3
493 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
494 }
495 } // namespace
496 } // aarch32
497 } // xnnpack
498
xnn_generate_qc8_igemm_fp32_ukernel_4x8__aarch32_neonv8_mlal_lane_ld64(xnn_code_buffer * code,size_t max_mr,size_t nc_mod_nr,size_t kc,size_t ks,const void * params)499 xnn_status_t xnn_generate_qc8_igemm_fp32_ukernel_4x8__aarch32_neonv8_mlal_lane_ld64(xnn_code_buffer* code, size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params) {
500 using namespace xnnpack::aarch32;
501 Generator g(code);
502 g.generate(false, max_mr, nc_mod_nr, kc, ks, nullptr);
503 g.finalize();
504 if (g.error() != xnnpack::Error::kNoError) {
505 return xnn_status_invalid_state;
506 }
507 return xnn_status_success;
508 }
509
xnn_generate_qc8_igemm_fp32_ukernel_4x8__aarch32_neonv8_mlal_lane_prfm_ld64(xnn_code_buffer * code,size_t max_mr,size_t nc_mod_nr,size_t kc,size_t ks,const void * params)510 xnn_status_t xnn_generate_qc8_igemm_fp32_ukernel_4x8__aarch32_neonv8_mlal_lane_prfm_ld64(xnn_code_buffer* code, size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params) {
511 using namespace xnnpack::aarch32;
512 Generator g(code);
513 g.generate(true, max_mr, nc_mod_nr, kc, ks, nullptr);
514 g.finalize();
515 if (g.error() != xnnpack::Error::kNoError) {
516 return xnn_status_invalid_state;
517 }
518 return xnn_status_success;
519 }
520