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(size_t max_mr, size_t nc_mod_nr, size_t kc, const void* params);
21 };
22
23
24 // void xnn_f32_igemm_minmax_ukernel_4x8__aarch32_neon_cortex_a53(
25 // size_t mr, r0
26 // size_t nc, r1
27 // size_t kc, r2 -> r5 -> sp + 68
28 // size_t ks, r3 -> sp + 72 -> r14
29 // const float**restrict a, sp + 112 -> (r5)
30 // const void*restrict w, sp + 116 -> r9
31 // uint8_t*restrict c, sp + 120 -> r11
32 // size_t cm_stride, sp + 124 -> (r6)
33 // size_t cn_stride, sp + 128 -> (r0)
34 // size_t a_offset, sp + 132 -> (r5)
35 // const float* zero, sp + 136 -> (r0)
36 // minmax_params*params, sp + 140 -> (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 // r0, r2 scratch temporaries for loads
43
44 // A0 r3 d0
45 // A1 r12 d1
46 // A2 r10 d2
47 // A3 r7 d3
48
49 // B r9 d8, d9, d10, d11
50 // B d12, d13, d14, d15
51
52 // C0 r11 d16-d17 q8 d18-d19 q9
53 // C1 r4 d20-d21 q10 d22-d23 q11
54 // C2 r8 d24-d25 q12 d26-d27 q13
55 // C3 r6 d28-d29 q14 d30-d31 q15
56
57 // Clamp (r5) d4 d5 d6 d7
58
59 // Converted from: src/f32-igemm/gen/4x8-minmax-aarch32-neon-cortex-a53.S
generate(size_t max_mr,size_t nc_mod_nr,size_t kc,const void * params)60 void Generator::generate(size_t max_mr, size_t nc_mod_nr, size_t kc, const void* params)
61 {
62 assert(nc_mod_nr < 8);
63 assert(kc != 0);
64 assert(kc % sizeof(float) == 0);
65
66 Label l0, l1, l2, l3, l4, l5, l6, l7, l8, l9, l10;
67
68 // Push 112 bytes
69 // r2 will be reloaded in outer loop. r3 is ks
70 push({r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, lr}); // +44
71 sub(sp, sp, 4); // 4
72 vpush({d8-d15}); // +64 = 112
73
74 ldr(r11, mem[sp, 120]); // c
75 ldr(r6, mem[sp, 124]); // cm_stride
76 ldr(r5, mem[sp, 112]); // a
77 ldr(r9, mem[sp, 116]); // w
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
92 align(8);
93 bind(l0);
94 // Load initial bias from w into accumulators
95 vldm(mem[r9]++, {d16-d19}); // Bias
96
97 vmov(q10, q8);
98 vmov(q11, q9);
99 vmov(q12, q8);
100 vmov(q13, q9);
101 pld(mem[r9, 0]); // Prefetch B
102 pld(mem[r9, 64]);
103 vmov(q14, q8);
104 pld(mem[r9, 128]);
105 pld(mem[r9, 192]);
106 vmov(q15, q9);
107 pld(mem[r9, 256]);
108 pld(mem[r9, 320]);
109
110 bind(l1);
111 // Load next 4 A pointers
112 ldr(r3, mem[r5, 0]);
113 ldr(r12, mem[r5, 4]);
114 ldr(r10, mem[r5, 8]);
115 ldr(r7, mem[r5, 12]);
116 add(r5, r5, 16);
117 pld(mem[r3, 0]); // Prefetch A
118 str(r5, mem[sp, 112]); // a
119 pld(mem[r3, 64]);
120 ldr(r0, mem[sp, 136]); // zero
121 pld(mem[r12, 0]);
122 ldr(r5, mem[sp, 132]); // a_offset
123 pld(mem[r12, 64]);
124 ldr(r2, mem[sp, 68]); // kc
125 pld(mem[r10, 0]);
126 pld(mem[r10, 64]);
127 pld(mem[r7, 0]);
128 pld(mem[r7, 64]);
129
130 // Add a_offset
131 cmp(r3, r0); // if a0 == zero
132 add(r3, r3, r5); // a0 += a_offset
133 moveq(r3, r0); // a0 = zero, else += a0 + a_offset
134 cmp(r12, r0); // if a1 == zero
135 add(r12, r12, r5); // a1 += a_offset
136 moveq(r12, r0); // a1 = zero, else += a1 + a_offset
137 cmp(r10, r0); // if a2 == zero
138 add(r10, r10, r5); // a2 += a_offset
139 moveq(r10, r0); // a2 = zero, else += a2 + a_offset
140 cmp(r7, r0); // if a3 == zero
141 add(r7, r7, r5); // a3 += a_offset
142 moveq(r7, r0); // a3 = zero, else += a3 + a_offset
143
144 subs(r5, r2, 16); // kc - 16
145 blo(l5); // less than 4 channels?
146
147 // Prologue
148 vld1_32({d0}, mem[r3]++); // A0
149 vld1_32({d1}, mem[r12]++); // A1
150 vld1_32({d2}, mem[r10]++); // A2
151 vld1_32({d3}, mem[r7]++); // A3
152 subs(r5, r5, 16);
153 vldm(mem[r9], {d8-d11}); // B0
154 ldr(r0, mem[r9, 56]); // B1 low VMOV is in BLOCK 0
155 ldr(r2, mem[r9, 60]); // B1 high
156 vldr(d13, mem[r9, 40]); // B1
157
158 blo(l3); // less than 4 channels? skip main loop
159
160 // Main loop - 4 floats of A (16 bytes)
161 // 32 FMA + 8 LD64 A + 8 LDR B
162 align(8);
163 bind(l2);
164 // First group of 16 FMA, Second group loads
165 // BLOCK 0
166 vld1_32({d4}, mem[r3]++); // A0
167 vmov(d15, r0, r2); // b1 VMOV b from second group
168 vmla_f32(q8, q4, d0[0]);
169 ldr(r0, mem[r12]); // A1 low
170 vmla_f32(q10, q4, d1[0]);
171 ldr(r2, mem[r12, 4]); // A1 high
172 vmla_f32(q12, q4, d2[0]);
173 pld(mem[r3, 128]); // Prefetch A0
174
175 // BLOCK 1
176 vldr(d12, mem[r9, 32]); // B1
177 vmov(d5, r0, r2); // a1 VMOV
178 vmla_f32(q14, q4, d3[0]);
179 ldr(r0, mem[r9, 72]); // B0 low
180 vmla_f32(q9, q5, d0[0]);
181 ldr(r2, mem[r9, 76]); // B0 high
182 vmla_f32(q11, q5, d1[0]);
183 pld(mem[r12, 128]); // Prefetch A1
184
185 // BLOCK 2
186 vld1_32({d6}, mem[r10]++); // A2
187 vmov(d9, r0, r2); // b0 VMOV
188 vmla_f32(q13, q5, d2[0]);
189 ldr(r0, mem[r7]); // A3 low
190 vmla_f32(q15, q5, d3[0]);
191 ldr(r2, mem[r7, 4]); // A3 high
192 vmla_f32(q8, q6, d0[1]);
193 pld(mem[r10, 128]); // Prefetch A2
194
195 // BLOCK 3
196 vldr(d14, mem[r9, 48]); // B1
197 vmov(d7, r0, r2); // a3 VMOV
198 vmla_f32(q10, q6, d1[1]);
199 ldr(r0, mem[r9, 88]); // B0 low
200 vmla_f32(q12, q6, d2[1]);
201 ldr(r2, mem[r9, 92]); // B0 high
202 vmla_f32(q14, q6, d3[1]);
203 pld(mem[r7, 128]); // Prefetch A3
204
205 // BLOCK 4
206 vldr(d8, mem[r9, 64]); // B0
207 vmov(d11, r0, r2); // B0 VMOV
208 vmla_f32(q9, q7, d0[1]);
209 ldr(r0, mem[r9, 104]); // B1 low VMOV is in BLOCK 0
210 vmla_f32(q11, q7, d1[1]);
211 ldr(r2, mem[r9, 108]); // B1 high
212 vmla_f32(q13, q7, d2[1]);
213 pld(mem[r9, 384]); // Prefetch B
214
215 // BLOCK 5
216 vldr(d10, mem[r9, 80]); // B0
217 vmov(d13, r0, r2); // b1 VMOV b from second group
218 vmla_f32(q15, q7, d3[1]);
219 ldr(r0, mem[r9, 120]); // B1 low VMOV is in BLOCK 0
220 nop();
221 ldr(r2, mem[r9, 124]); // B1 high
222 nop();
223 pld(mem[r9, 448]); // Prefetch B
224
225 // Second group of 16 FMA, First group of loads
226 // BLOCK 0
227 vld1_32({d0}, mem[r3]++); // A0
228 vmov(d15, r0, r2); // b1 VMOV b from second group
229 vmla_f32(q8, q4, d4[0]);
230 ldr(r0, mem[r12, 8]); // A1 low
231 vmla_f32(q10, q4, d5[0]);
232 ldr(r2, mem[r12, 12]); // A1 high
233 vmla_f32(q12, q4, d6[0]);
234 // NOP
235
236 // BLOCK 1
237 vldr(d12, mem[r9, 96]); // B1
238 vmov(d1, r0, r2); // a1 VMOV
239 vmla_f32(q14, q4, d7[0]);
240 ldr(r0, mem[r9, 136]); // B0 low
241 vmla_f32(q9, q5, d4[0]);
242 ldr(r2, mem[r9, 140]); // B0 high
243 vmla_f32(q11, q5, d5[0]);
244 // NOP
245
246 // BLOCK 2
247 vld1_32({d2}, mem[r10]++); // A2
248 vmov(d9, r0, r2); // b0 VMOV
249 vmla_f32(q13, q5, d6[0]);
250 ldr(r0, mem[r7, 8]); // A3 low
251 vmla_f32(q15, q5, d7[0]);
252 ldr(r2, mem[r7, 12]); // A3 high
253 vmla_f32(q8, q6, d4[1]);
254 // NOP
255
256 // BLOCK 3
257 vldr(d14, mem[r9, 112]); // B1
258 vmov(d3, r0, r2); // a3 VMOV
259 vmla_f32(q10, q6, d5[1]);
260 ldr(r0, mem[r9, 152]); // B0 low
261 vmla_f32(q12, q6, d6[1]);
262 ldr(r2, mem[r9, 156]); // B0 high
263 vmla_f32(q14, q6, d7[1]);
264 add(r12, r12, 16); // A1++
265
266 // BLOCK 4
267 vldr(d8, mem[r9, 128]); // B0
268 vmov(d11, r0, r2); // B0 VMOV
269 vmla_f32(q9, q7, d4[1]);
270 ldr(r0, mem[r9, 168]); // B1 low
271 vmla_f32(q11, q7, d5[1]);
272 ldr(r2, mem[r9, 172]); // B1 high
273 vmla_f32(q13, q7, d6[1]);
274 add(r7, r7, 16); // A3++
275
276 // BLOCK 5
277 vldr(d10, mem[r9, 144]); // B0
278 vmov(d13, r0, r2); // b1 VMOV b
279 vmla_f32(q15, q7, d7[1]);
280 ldr(r0, mem[r9, 184]); // B1 low VMOV is in BLOCK 0
281 subs(r5, r5, 16);
282 ldr(r2, mem[r9, 188]); // B1 high
283 add(r9, r9, 128); // B++
284 bhs(l2);
285
286 // Epilogue - 4 floats of A (16 bytes)
287 bind(l3);
288 // First group of 16 FMA, Second group loads
289 // BLOCK 0
290 vld1_32({d4}, mem[r3]++); // A0
291 vmov(d15, r0, r2); // b1 VMOV b from second group
292 vmla_f32(q8, q4, d0[0]);
293 ldr(r0, mem[r12]); // A1 low
294 vmla_f32(q10, q4, d1[0]);
295 ldr(r2, mem[r12, 4]); // A1 high
296 vmla_f32(q12, q4, d2[0]);
297 // NOP
298
299 // BLOCK 1
300 vldr(d12, mem[r9, 32]); // B1
301 vmov(d5, r0, r2); // a1 VMOV
302 vmla_f32(q14, q4, d3[0]);
303 ldr(r0, mem[r9, 72]); // B0 low
304 vmla_f32(q9, q5, d0[0]);
305 ldr(r2, mem[r9, 76]); // B0 high
306 vmla_f32(q11, q5, d1[0]);
307 // NOP
308
309 // BLOCK 2
310 vld1_32({d6}, mem[r10]++); // A2
311 vmov(d9, r0, r2); // b0 VMOV
312 vmla_f32(q13, q5, d2[0]);
313 ldr(r0, mem[r7]); // A3 low
314 vmla_f32(q15, q5, d3[0]);
315 ldr(r2, mem[r7, 4]); // A3 high
316 vmla_f32(q8, q6, d0[1]);
317 // NOP
318
319 // BLOCK 3
320 vldr(d14, mem[r9, 48]); // B1
321 vmov(d7, r0, r2); // a3 VMOV
322 vmla_f32(q10, q6, d1[1]);
323 ldr(r0, mem[r9, 88]); // B0 low
324 vmla_f32(q12, q6, d2[1]);
325 ldr(r2, mem[r9, 92]); // B0 high
326 vmla_f32(q14, q6, d3[1]);
327 // NOP
328
329 // BLOCK 4
330 vldr(d8, mem[r9, 64]); // B0
331 vmov(d11, r0, r2); // B0 VMOV
332 vmla_f32(q9, q7, d0[1]);
333 ldr(r0, mem[r9, 104]); // B1 low
334 vmla_f32(q11, q7, d1[1]);
335 ldr(r2, mem[r9, 108]); // B1 high
336 vmla_f32(q13, q7, d2[1]);
337 // NOP
338
339 // BLOCK 5
340 vldr(d10, mem[r9, 80]); // B0
341 vmov(d13, r0, r2); // b1 VMOV b
342 vmla_f32(q15, q7, d3[1]);
343 ldr(r0, mem[r9, 120]); // B1 low VMOV is in BLOCK 0
344 nop();
345 ldr(r2, mem[r9, 124]); // B1 high
346 nop();
347 nop();
348
349 // Second group of 16 FMA, First group of loads
350 // BLOCK 0
351 vldr(d12, mem[r9, 96]); // B1
352 vmov(d15, r0, r2); // b1 VMOV b from second group
353 vmla_f32(q8, q4, d4[0]);
354 vmla_f32(q10, q4, d5[0]);
355 vmla_f32(q12, q4, d6[0]);
356
357 // BLOCK 1
358 vldr(d14, mem[r9, 112]); // B1
359 vmla_f32(q14, q4, d7[0]);
360 vmla_f32(q9, q5, d4[0]);
361 vmla_f32(q11, q5, d5[0]);
362 add(r12, r12, 8); // A1++
363
364 // BLOCK 2
365 add(r7, r7, 8); // A3++ VLDR B1 land_s here
366 add(r9, r9, 128); // B++
367 vmla_f32(q13, q5, d6[0]);
368 vmla_f32(q15, q5, d7[0]);
369 vmla_f32(q8, q6, d4[1]);
370
371 // BLOCK 3
372 vmla_f32(q10, q6, d5[1]);
373 vmla_f32(q12, q6, d6[1]);
374 vmla_f32(q14, q6, d7[1]);
375 tst(r5, 15);
376
377 // BLOCK 4
378 vmla_f32(q9, q7, d4[1]);
379 vmla_f32(q11, q7, d5[1]);
380 vmla_f32(q13, q7, d6[1]);
381
382 // BLOCK 5
383 vmla_f32(q15, q7, d7[1]);
384
385 // Is there a remainder?- 1 to 3 floats of A (4, 8 or 12 bytes)
386 bne(l5);
387
388 align(8);
389 bind(l4);
390 ldr(r5, mem[sp, 112]); // a
391 subs(r14, r14, 16); // ks -= MR * sizeof(void*)
392
393 // ks loop
394 bhi(l1);
395
396 // Load params pointer
397 ldr(r0, mem[sp, 128]); // cn_stride
398 ldr(r2, mem[sp, 140]); // params
399 ldr(r14, mem[sp, 72]); // p = ks
400 subs(r1, r1, 8);
401
402 // Load min/max values
403 vld1r_32({d4,d5}, mem[r2]++);
404 vld1r_32({d6,d7}, mem[r2]);
405
406 // Clamp
407 vmax_f32(q8, q8, q2);
408 vmax_f32(q9, q9, q2);
409 vmax_f32(q10, q10, q2);
410 vmax_f32(q11, q11, q2);
411 vmax_f32(q12, q12, q2);
412 vmax_f32(q13, q13, q2);
413 vmax_f32(q14, q14, q2);
414 vmax_f32(q15, q15, q2);
415 vmin_f32(q8, q8, q3);
416 vmin_f32(q9, q9, q3);
417 vmin_f32(q10, q10, q3);
418 vmin_f32(q11, q11, q3);
419 vmin_f32(q12, q12, q3);
420 vmin_f32(q13, q13, q3);
421 vmin_f32(q14, q14, q3);
422 vmin_f32(q15, q15, q3);
423
424 // Store full 4 x 8
425 blo(l7);
426 vst1_32({d28-d31}, mem[r6], r0);
427 vst1_32({d24-d27}, mem[r8], r0);
428 vst1_32({d20-d23}, mem[r4], r0);
429 vst1_32({d16-d19}, mem[r11], r0);
430
431 sub(r5, r5, r14); // a -= ks
432
433 bhi(l0);
434
435 vpop({d8-d15});
436 add(sp, sp, 12); // skip pad, r2, r3
437 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
438
439 align(8);
440 bind(l5);
441 // Is there a remainder?- 2 floats of A (8 bytes)
442 tst(r5, 8);
443 beq(l6);
444
445 // Remainder - 2 floats of A (8 bytes)
446 vld1_32({d0}, mem[r3]++); // A0
447 vldm(mem[r9]++, {d8-d11}); // B0
448 vld1_32({d1}, mem[r12]++); // A1
449 vld1_32({d2}, mem[r10]++); // A2
450 vld1_32({d3}, mem[r7]++); // A3
451
452 vmla_f32(q8, q4, d0[0]);
453 vmla_f32(q9, q5, d0[0]);
454 vmla_f32(q10, q4, d1[0]);
455 vmla_f32(q11, q5, d1[0]);
456 vldm(mem[r9]++, {d12-d15}); // B1
457 vmla_f32(q12, q4, d2[0]);
458 vmla_f32(q13, q5, d2[0]);
459 vmla_f32(q14, q4, d3[0]);
460 vmla_f32(q15, q5, d3[0]);
461 vmla_f32(q8, q6, d0[1]);
462 vmla_f32(q9, q7, d0[1]);
463 vmla_f32(q10, q6, d1[1]);
464 vmla_f32(q11, q7, d1[1]);
465 vmla_f32(q12, q6, d2[1]);
466 vmla_f32(q13, q7, d2[1]);
467 vmla_f32(q14, q6, d3[1]);
468 vmla_f32(q15, q7, d3[1]);
469
470 // Is there a remainder?- 1 float of A (4 bytes)
471 tst(r5, 4);
472 beq(l4);
473
474 bind(l6);
475 // Remainder- 1 float of A (4 bytes)
476 vldm(mem[r3]++, {s0}); // A0
477 vldm(mem[r9]++, {d8-d11}); // B0
478 vldm(mem[r12]++, {s2}); // A1
479 vldm(mem[r10]++, {s4}); // A2
480 vldm(mem[r7]++, {s6}); // A3
481 vmla_f32(q8, q4, d0[0]);
482 vmla_f32(q9, q5, d0[0]);
483 vmla_f32(q10, q4, d1[0]);
484 vmla_f32(q11, q5, d1[0]);
485 vmla_f32(q12, q4, d2[0]);
486 vmla_f32(q13, q5, d2[0]);
487 vmla_f32(q14, q4, d3[0]);
488 vmla_f32(q15, q5, d3[0]);
489 b(l4);
490
491 // Store odd width
492 bind(l7);
493 tst(r1, 4);
494 beq(l8);
495 vst1_32({d28-d29}, mem[r6]++);
496 vst1_32({d24-d25}, mem[r8]++);
497 vmov(q14, q15);
498 vmov(q12, q13);
499 vst1_32({d20-d21}, mem[r4]++);
500 vst1_32({d16-d17}, mem[r11]++);
501 vmov(q10, q11);
502 vmov(q8, q9);
503
504 bind(l8);
505 tst(r1, 2);
506 beq(l9);
507 vst1_32({d28}, mem[r6]++);
508 vst1_32({d24}, mem[r8]++);
509 vmov(d28, d29);
510 vmov(d24, d25);
511 vst1_32({d20}, mem[r4]++);
512 vst1_32({d16}, mem[r11]++);
513 vmov(d20, d21);
514 vmov(d16, d17);
515
516 bind(l9);
517 tst(r1, 1);
518 beq(l10);
519 vst1_32({d28[0]}, mem[r6]++);
520 vst1_32({d24[0]}, mem[r8]++);
521 vst1_32({d20[0]}, mem[r4]++);
522 vst1_32({d16[0]}, mem[r11]++);
523
524 bind(l10);
525 vpop({d8-d15});
526 add(sp, sp, 12); // skip pad, r2, r3
527 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
528 }
529 } // namespace
530 } // aarch32
531 } // xnnpack
532
xnn_generate_f32_igemm_ukernel_4x8__aarch32_neon_cortex_a53(xnn_code_buffer * code,size_t max_mr,size_t nc_mod_nr,size_t kc,size_t ks,const void * params)533 xnn_status_t xnn_generate_f32_igemm_ukernel_4x8__aarch32_neon_cortex_a53(xnn_code_buffer* code, size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params) {
534 using namespace xnnpack::aarch32;
535 Generator g(code);
536 assert(params != nullptr);
537 g.generate(max_mr, nc_mod_nr, kc, nullptr);
538 g.finalize();
539 if (g.error() != xnnpack::Error::kNoError) {
540 return xnn_status_invalid_state;
541 }
542 return xnn_status_success;
543 }
544