xref: /aosp_15_r20/external/XNNPACK/src/f32-gemm/4x8-aarch32-neon-cortex-a7.cc (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
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/gemm.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_gemm_minmax_ukernel_4x8__aarch32_neon_cortex_a7(
25 //     size_t mr,                            r0
26 //     size_t nc,                            r1
27 //     size_t kc,                            r2 -> r5
28 //     const uint8_t*restrict a,             r3
29 //     size_t a_stride,          sp + 96  -> (r7)
30 //     const void*restrict w,    sp + 100 -> r9
31 //     uint8_t*restrict c,       sp + 104 -> r11
32 //     size_t cm_stride,         sp + 108 -> (r6)
33 //     size_t cn_stride,         sp + 112 -> r7
34 //     const union xnn_f32_minmax_params params)  sp + 116 -> (r7)
35 
36 // d8-d15, r4-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.
37 
38 // Register usage
39 
40 // A0   r3  d0
41 // A1  r12  d1
42 // A2  r10  d2
43 // A3   r0  d3
44 
45 // B    r9  d8,  d9, d10, d11
46 // B       d12, d13, d14, d15
47 
48 // C0  r11 d16-d17  q8  d18-d19  q9
49 // C1   r4 d20-d21 q10  d22-d23 q11
50 // C2   r8 d24-d25 q12  d26-d27 q13
51 // C3   r6 d28-d29 q14  d30-d31 q15
52 
53 // Clamp (r5) d4 d5 d6 d7
54 
55 // Converted from: src/f32-gemm/gen/4x8-minmax-aarch32-neon-cortex-a7.S
generate(size_t max_mr,size_t nc_mod_nr,size_t kc,const void * params)56 void Generator::generate(size_t max_mr, size_t nc_mod_nr, size_t kc, const void* params)
57 {
58   assert(nc_mod_nr < 8);
59   assert(kc != 0);
60   assert(kc % sizeof(float) == 0);
61 
62   Label l0, l1, l2, l3, l4, l5, l6, l7;
63 
64   // Push 96 bytes
65   push({r4, r5, r6, r7, r8, r9, r10, r11}); // 32
66   vpush({d8-d15}); // +64 = 96
67 
68   ldr(r7, mem[sp, 96]); // a_stride
69   ldr(r11, mem[sp, 104]); // c
70   ldr(r6, mem[sp, 108]); // cm_stride
71   ldr(r9, mem[sp, 100]); // w
72   ldr(r5, mem[sp, 116]); // params
73 
74   // Clamp A and C pointers
75   cmp(r0, 2); // if mr >= 2
76   add(r12, r3, r7); //   a1 = a0 + a_stride
77   add(r4, r11, r6); //   c1 = c0 + cm_stride
78   movlo(r12, r3); // a1
79   movlo(r4, r11); // c1
80   // if mr > 2
81   add(r10, r12, r7); //   a2 = a1 + a_stride
82   add(r8, r4, r6); //   c2 = c1 + cm_stride
83   movls(r10, r12); // a2
84   movls(r8, r4); // c2
85 
86   cmp(r0, 4); // if mr >=4
87   add(r0, r10, r7); //   a3 = a2 + a_stride
88   add(r6, r8, r6); //   c3 = c2 + cm_stride
89   movlo(r0, r10); // a3
90   movlo(r6, r8); // c3
91 
92   // Load min/max values
93   vld1r_32({d4, d5}, mem[r5]++);
94   ldr(r7, mem[sp, 112]); // cn_stride
95   vld1r_32({d6, d7}, mem[r5]);
96 
97   bind(l0);
98   // Load initial bias from w into accumulators
99   vldm(mem[r9]++, {d16-d19}); // Bias
100   subs(r5, r2, 8);
101   vmov(q10, q8);
102   vmov(q11, q9);
103   vmov(q12, q8);
104   vmov(q13, q9);
105   vmov(q14, q8);
106   vmov(q15, q9);
107 
108   pld(mem[r3, 0]); // Prefetch A
109   pld(mem[r3, 64]);
110   pld(mem[r12, 0]);
111   pld(mem[r12, 64]);
112   pld(mem[r10, 0]);
113   pld(mem[r10, 64]);
114   pld(mem[r0, 0]);
115   pld(mem[r0, 64]);
116   pld(mem[r9, 0]); // Prefetch B
117   pld(mem[r9, 64]);
118   pld(mem[r9, 128]);
119   pld(mem[r9, 192]);
120   pld(mem[r9, 256]);
121   pld(mem[r9, 320]);
122   pld(mem[r9, 384]);
123   pld(mem[r9, 448]);
124   blo(l3); // less than 2 channels?
125 
126   // Main loop - 2 floats of A (8 bytes)
127   bind(l1);
128   vld1_32({d0}, mem[r3]++); // A0
129   vldm(mem[r9]++, {d8-d11}); // B0
130   vld1_32({d1}, mem[r12]++); // A1
131   vld1_32({d2}, mem[r10]++); // A2
132   vld1_32({d3}, mem[r0]++); // A3
133   vldm(mem[r9]++, {d12-d15}); // B1
134   vmla_f32(q8, q4, d0[0]);
135   vmla_f32(q9, q5, d0[0]);
136   vmla_f32(q10, q4, d1[0]);
137   vmla_f32(q13, q5, d2[0]);
138   vmla_f32(q11, q5, d1[0]);
139   vmla_f32(q12, q4, d2[0]);
140   vmla_f32(q14, q4, d3[0]);
141   vmla_f32(q15, q5, d3[0]);
142   vmla_f32(q8, q6, d0[1]);
143   vmla_f32(q9, q7, d0[1]);
144   vmla_f32(q10, q6, d1[1]);
145   vmla_f32(q11, q7, d1[1]);
146   subs(r5, r5, 8);
147   vmla_f32(q12, q6, d2[1]);
148   vmla_f32(q13, q7, d2[1]);
149   vmla_f32(q14, q6, d3[1]);
150   vmla_f32(q15, q7, d3[1]);
151   pld(mem[r9, 448]); // Prefetch B
152   pld(mem[r3, 128]); // Prefetch A0
153   pld(mem[r12, 128]); // Prefetch A1
154   pld(mem[r10, 128]); // Prefetch A2
155   pld(mem[r0, 128]); // Prefetch A3
156   bhs(l1);
157 
158   // Is there a remainder?- 1 float of A (4 bytes)
159   tst(r5, 4);
160   bne(l3);
161 
162   bind(l2);
163   // Clamp
164   vmax_f32(q8, q8, q2);
165   subs(r1, r1, 8);
166   vmax_f32(q9, q9, q2);
167   vmax_f32(q10, q10, q2);
168   vmax_f32(q11, q11, q2);
169   vmax_f32(q12, q12, q2);
170   vmax_f32(q13, q13, q2);
171   vmax_f32(q14, q14, q2);
172   vmax_f32(q15, q15, q2);
173   vmin_f32(q8, q8, q3);
174   vmin_f32(q9, q9, q3);
175   vmin_f32(q10, q10, q3);
176   vmin_f32(q11, q11, q3);
177   vmin_f32(q12, q12, q3);
178   vmin_f32(q13, q13, q3);
179   vmin_f32(q14, q14, q3);
180   vmin_f32(q15, q15, q3);
181 
182   // Store full 4 x 8
183   blo(l4);
184   vst1_32({d16-d19}, mem[r11], r7);
185   sub(r0, r0, r2);
186   vst1_32({d20-d23}, mem[r4], r7);
187   sub(r10, r10, r2);
188   vst1_32({d24-d27}, mem[r8], r7);
189   sub(r12, r12, r2);
190   vst1_32({d28-d31}, mem[r6], r7);
191   sub(r3, r3, r2);
192   bhi(l0);
193 
194   vpop({d8-d15});
195   pop({r4, r5, r6, r7, r8, r9, r10, r11});
196   bx(lr);
197 
198   bind(l3);
199   // Remainder- 1 float of A (4 bytes)
200   vldm(mem[r3]++, {s0}); // A0
201   vldm(mem[r9]++, {d8-d11}); // B0
202   vldm(mem[r12]++, {s2}); // A1
203   vldm(mem[r10]++, {s4}); // A2
204   vldm(mem[r0]++, {s6}); // A3
205   vmla_f32(q8, q4, d0[0]);
206   vmla_f32(q9, q5, d0[0]);
207   vmla_f32(q10, q4, d1[0]);
208   vmla_f32(q11, q5, d1[0]);
209   vmla_f32(q12, q4, d2[0]);
210   vmla_f32(q13, q5, d2[0]);
211   vmla_f32(q14, q4, d3[0]);
212   vmla_f32(q15, q5, d3[0]);
213   b(l2);
214 
215   // Store odd width
216   bind(l4);
217   tst(r1, 4);
218   beq(l5);
219   vst1_32({d16-d17}, mem[r11]++);
220   vst1_32({d20-d21}, mem[r4]++);
221   vmov(q8, q9);
222   vmov(q10, q11);
223   vst1_32({d24-d25}, mem[r8]++);
224   vst1_32({d28-d29}, mem[r6]++);
225   vmov(q12, q13);
226   vmov(q14, q15);
227 
228   bind(l5);
229   tst(r1, 2);
230   beq(l6);
231   vst1_32({d16}, mem[r11]++);
232   vst1_32({d20}, mem[r4]++);
233   vmov(d16, d17);
234   vmov(d20, d21);
235   vst1_32({d24}, mem[r8]++);
236   vst1_32({d28}, mem[r6]++);
237   vmov(d24, d25);
238   vmov(d28, d29);
239 
240   bind(l6);
241   tst(r1, 1);
242   beq(l7);
243   vst1_32({d16[0]}, mem[r11]);
244   vst1_32({d20[0]}, mem[r4]);
245   vst1_32({d24[0]}, mem[r8]);
246   vst1_32({d28[0]}, mem[r6]);
247 
248   bind(l7);
249   vpop({d8-d15});
250   pop({r4, r5, r6, r7, r8, r9, r10, r11});
251   bx(lr);
252 }
253 }  // namespace
254 }  // aarch32
255 }  // xnnpack
256 
xnn_generate_f32_gemm_ukernel_4x8__aarch32_neon_cortex_a7(xnn_code_buffer * code,size_t max_mr,size_t nc_mod_nr,size_t kc,const void * params)257 xnn_status_t xnn_generate_f32_gemm_ukernel_4x8__aarch32_neon_cortex_a7(xnn_code_buffer* code, size_t max_mr, size_t nc_mod_nr, size_t kc, const void* params) {
258   using namespace xnnpack::aarch32;
259   Generator g(code);
260   assert(params != nullptr);
261   g.generate(max_mr, nc_mod_nr, kc, nullptr);
262   g.finalize();
263   if (g.error() != xnnpack::Error::kNoError) {
264     return xnn_status_invalid_state;
265   }
266   return xnn_status_success;
267 }
268