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