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