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