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_a7(
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 -> r2
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 -> (r7)
34 // size_t a_offset, sp + 132 -> (r5)
35 // const float* zero, sp + 136 -> (r7)
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 // A0 r3 d0
43 // A1 r12 d1
44 // A2 r10 d2
45 // A3 r0 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-igemm/gen/4x8-minmax-aarch32-neon-cortex-a7.S
generate(size_t max_mr,size_t nc_mod_nr,size_t kc,const void * params)58 void Generator::generate(size_t max_mr, size_t nc_mod_nr, size_t kc, const void* params)
59 {
60 assert(nc_mod_nr < 8);
61 assert(kc != 0);
62 assert(kc % sizeof(float) == 0);
63
64 Label l0, l1, l2, l3, l4, l5, l6, l7, l8;
65
66 // Push 112 bytes
67 // r2 will be reloaded in outer loop. r3 is ks
68 push({r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, lr}); // +44
69 sub(sp, sp, 4); // 4
70 vpush({d8-d15}); // +64 = 112
71
72 ldr(r11, mem[sp, 120]); // c
73 ldr(r6, mem[sp, 124]); // cm_stride
74 ldr(r2, mem[sp, 112]); // a
75 ldr(r9, mem[sp, 116]); // w
76 ldr(r5, mem[sp, 140]); // params
77 mov(r14, r3); // p = ks
78
79 // Clamp C pointers
80 cmp(r0, 2); // if mr >= 2
81 add(r4, r11, r6); // c1 = c0 + cm_stride
82 movlo(r4, r11); // c1
83 // if mr > 2
84 add(r8, r4, r6); // c2 = c1 + cm_stride
85 movls(r8, r4); // c2
86 cmp(r0, 4); // if mr >=4
87 add(r6, r8, r6); // c3 = c2 + cm_stride
88 movlo(r6, r8); // c3
89
90 // Load min/max values
91 vld1r_32({d4, d5}, mem[r5]++);
92 vld1r_32({d6, d7}, mem[r5]);
93
94 bind(l0);
95 // Load initial bias from w into accumulators
96 vldm(mem[r9]++, {d16-d19}); // Bias
97 vmov(q10, q8);
98 vmov(q11, q9);
99 vmov(q12, q8);
100 vmov(q13, q9);
101 vmov(q14, q8);
102 vmov(q15, q9);
103
104 pld(mem[r9, 0]); // Prefetch B
105 pld(mem[r9, 64]);
106 pld(mem[r9, 128]);
107 pld(mem[r9, 192]);
108 pld(mem[r9, 256]);
109 pld(mem[r9, 320]);
110 pld(mem[r9, 384]);
111 pld(mem[r9, 448]);
112 bind(l1);
113 // Load next 4 A pointers
114 ldr(r3, mem[r2, 0]);
115 ldr(r12, mem[r2, 4]);
116 ldr(r10, mem[r2, 8]);
117 ldr(r0, mem[r2, 12]);
118 add(r2, r2, 16);
119
120 // Add a_offset
121 ldr(r5, mem[sp, 132]); // a_offset
122 ldr(r7, mem[sp, 136]); // zero
123 cmp(r3, r7); // if a0 == zero
124 add(r3, r3, r5); // a0 += a_offset
125 moveq(r3, r7); // a0 = zero, else += a0 + a_offset
126 cmp(r12, r7); // if a1 == zero
127 add(r12, r12, r5); // a1 += a_offset
128 moveq(r12, r7); // a1 = zero, else += a1 + a_offset
129 cmp(r10, r7); // if a2 == zero
130 add(r10, r10, r5); // a2 += a_offset
131 moveq(r10, r7); // a2 = zero, else += a2 + a_offset
132 cmp(r0, r7); // if a3 == zero
133 add(r0, r0, r5); // a3 += a_offset
134 ldr(r5, mem[sp, 68]); // kc
135 moveq(r0, r7); // a3 = zero, else += a3 + a_offset
136
137 pld(mem[r3, 0]); // Prefetch A
138 pld(mem[r3, 64]);
139 pld(mem[r12, 0]);
140 pld(mem[r12, 64]);
141 pld(mem[r10, 0]);
142 pld(mem[r10, 64]);
143 pld(mem[r0, 0]);
144 pld(mem[r0, 64]);
145
146 subs(r5, r5, 8); // kc - 8
147 blo(l4); // less than 2 channels?
148
149 // Main loop - 2 floats of A (8 bytes)
150 bind(l2);
151 vld1_32({d0}, mem[r3]++); // A0
152 vldm(mem[r9]++, {d8-d11}); // B0
153 vld1_32({d1}, mem[r12]++); // A1
154 vld1_32({d2}, mem[r10]++); // A2
155 vld1_32({d3}, mem[r0]++); // A3
156 vldm(mem[r9]++, {d12-d15}); // B1
157
158 vmla_f32(q8, q4, d0[0]);
159 vmla_f32(q9, q5, d0[0]);
160 vmla_f32(q10, q4, d1[0]);
161 vmla_f32(q11, q5, d1[0]);
162 vmla_f32(q12, q4, d2[0]);
163 vmla_f32(q13, q5, d2[0]);
164 vmla_f32(q14, q4, d3[0]);
165 vmla_f32(q15, q5, d3[0]);
166 vmla_f32(q8, q6, d0[1]);
167 vmla_f32(q9, q7, d0[1]);
168 vmla_f32(q10, q6, d1[1]);
169 vmla_f32(q11, q7, d1[1]);
170 subs(r5, r5, 8);
171 vmla_f32(q12, q6, d2[1]);
172 vmla_f32(q13, q7, d2[1]);
173 vmla_f32(q14, q6, d3[1]);
174 vmla_f32(q15, q7, d3[1]);
175 pld(mem[r9, 448]); // Prefetch B
176 pld(mem[r3, 128]); // Prefetch A0
177 pld(mem[r12, 128]); // Prefetch A1
178 pld(mem[r10, 128]); // Prefetch A2
179 pld(mem[r0, 128]); // Prefetch A3
180 bhs(l2);
181
182 // Is there a remainder?- 1 float of A (4 bytes)
183 tst(r5, 4);
184 bne(l4);
185
186 bind(l3);
187 // ks loop
188 subs(r14, r14, 16); // ks -= MR * sizeof(void*)
189 bhi(l1);
190
191 ldr(r7, mem[sp, 128]); // cn_stride
192 ldr(r14, mem[sp, 72]); // p = ks
193
194 // Clamp
195 vmax_f32(q8, q8, q2);
196 subs(r1, r1, 8);
197 vmax_f32(q9, q9, q2);
198 vmax_f32(q10, q10, q2);
199 vmax_f32(q11, q11, q2);
200 vmax_f32(q12, q12, q2);
201 vmax_f32(q13, q13, q2);
202 vmax_f32(q14, q14, q2);
203 vmax_f32(q15, q15, q2);
204 vmin_f32(q8, q8, q3);
205 vmin_f32(q9, q9, q3);
206 vmin_f32(q10, q10, q3);
207 vmin_f32(q11, q11, q3);
208 vmin_f32(q12, q12, q3);
209 vmin_f32(q13, q13, q3);
210 vmin_f32(q14, q14, q3);
211 vmin_f32(q15, q15, q3);
212
213 // Store full 4 x 8
214 blo(l5);
215 vst1_32({d28-d31}, mem[r6], r7);
216 vst1_32({d24-d27}, mem[r8], r7);
217 vst1_32({d20-d23}, mem[r4], r7);
218 vst1_32({d16-d19}, mem[r11], r7);
219 sub(r2, r2, r14); // a -= ks
220 bhi(l0);
221
222 vpop({d8-d15});
223 add(sp, sp, 12); // skip pad, r2, r3
224 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
225
226 bind(l4);
227 // Remainder- 1 float of A (4 bytes)
228 vldm(mem[r3]++, {s0}); // A0
229 vldm(mem[r9]++, {d8-d11}); // B0
230 vldm(mem[r12]++, {s2}); // A1
231 vldm(mem[r10]++, {s4}); // A2
232 vldm(mem[r0]++, {s6}); // A3
233 vmla_f32(q8, q4, d0[0]);
234 vmla_f32(q9, q5, d0[0]);
235 vmla_f32(q10, q4, d1[0]);
236 vmla_f32(q11, q5, d1[0]);
237 vmla_f32(q12, q4, d2[0]);
238 vmla_f32(q13, q5, d2[0]);
239 vmla_f32(q14, q4, d3[0]);
240 vmla_f32(q15, q5, d3[0]);
241 b(l3);
242
243 // Store odd width
244 bind(l5);
245 tst(r1, 4);
246 beq(l6);
247 vst1_32({d28-d29}, mem[r6]++);
248 vst1_32({d24-d25}, mem[r8]++);
249 vmov(q14, q15);
250 vmov(q12, q13);
251 vst1_32({d20-d21}, mem[r4]++);
252 vst1_32({d16-d17}, mem[r11]++);
253 vmov(q10, q11);
254 vmov(q8, q9);
255
256 bind(l6);
257 tst(r1, 2);
258 beq(l7);
259 vst1_32({d28}, mem[r6]++);
260 vst1_32({d24}, mem[r8]++);
261 vmov(d28, d29);
262 vmov(d24, d25);
263 vst1_32({d20}, mem[r4]++);
264 vst1_32({d16}, mem[r11]++);
265 vmov(d20, d21);
266 vmov(d16, d17);
267
268 bind(l7);
269 tst(r1, 1);
270 beq(l8);
271 vst1_32({d28[0]}, mem[r6]++);
272 vst1_32({d24[0]}, mem[r8]++);
273 vst1_32({d20[0]}, mem[r4]++);
274 vst1_32({d16[0]}, mem[r11]++);
275
276 bind(l8);
277 vpop({d8-d15});
278 add(sp, sp, 12); // skip pad, r2, r3
279 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
280 }
281 } // namespace
282 } // aarch32
283 } // xnnpack
284
xnn_generate_f32_igemm_ukernel_4x8__aarch32_neon_cortex_a7(xnn_code_buffer * code,size_t max_mr,size_t nc_mod_nr,size_t kc,size_t ks,const void * params)285 xnn_status_t xnn_generate_f32_igemm_ukernel_4x8__aarch32_neon_cortex_a7(xnn_code_buffer* code, size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params) {
286 using namespace xnnpack::aarch32;
287 Generator g(code);
288 assert(params != nullptr);
289 g.generate(max_mr, nc_mod_nr, kc, nullptr);
290 g.finalize();
291 if (g.error() != xnnpack::Error::kNoError) {
292 return xnn_status_invalid_state;
293 }
294 return xnn_status_success;
295 }
296