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, size_t ks, const void* params);
21 };
22
23
24 // void xnn_qs8_igemm_minmax_rndnu_ukernel_4x8c4__aarch32_neondot_ld64(
25 // size_t mr, r0
26 // size_t nc, r1
27 // size_t kc, r2 -> r5 -> sp + 52
28 // size_t ks, r3 -> sp + 56 -> r14
29 // const int8_t**restrict a, sp + 96 -> r2
30 // const void*restrict w, sp + 100 -> r9
31 // int8_t*restrict c, sp + 104 -> r11
32 // size_t cm_stride, sp + 108 -> (r6)
33 // size_t cn_stride, sp + 112 -> (r7)
34 // size_t a_offset, sp + 116 -> (r5)
35 // const int8_t* zero, sp + 120 -> (r7)
36 // xnn_qs8_conv_minmax_params*params); sp + 124 -> (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 q2 q3 q4 q5
48
49 // C0 r11 d16-d17 q8 d18-d19 q9
50 // C1 r4 d20-d21 q10 d22-d23 q11
51 // C2 r8 d24-d25 q12 d26-d27 q13
52 // C3 r6 d28-d29 q14 d30-d31 q15
53
54 // unused q7
55
56 // params structure is 16 bytes
57 // struct {
58 // int32_t right_pre_shift; d12[0]
59 // int32_t multiplier; d12[1]
60 // int32_t right_post_shift; d13[0]
61 // int16_t output_zero_point; d13[2]
62 // int8_t output_min; d13[6]
63 // int8_t output_max; d13[7]
64 // } rndnu_neon;
65
66 // Converted from: src/qs8-igemm/gen/4x8c4-minmax-rndnu-aarch32-neondot-ld64.S
generate(size_t max_mr,size_t nc_mod_nr,size_t kc,size_t ks,const void * params)67 void Generator::generate(size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params)
68 {
69 assert(nc_mod_nr < 8);
70 assert(kc != 0);
71 assert(ks != 0);
72
73 Label l0, l1, l2, l3, l4, l5, l6, l7, l8;
74
75 add(r2, r2, 3); // kc = (kc + 3) & ~3
76 bic(r2, r2, 3);
77
78 // Push 96 bytes
79 // r2 will be reloaded in outer loop. r3 is ks
80 push({r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, lr}); // +44
81 sub(sp, sp, 4); // 4
82 vpush({d8-d13}); // +48 = 96
83
84 ldr(r11, mem[sp, 104]); // c
85 ldr(r6, mem[sp, 108]); // cm_stride
86 ldr(r2, mem[sp, 96]); // a
87 ldr(r9, mem[sp, 100]); // w
88 ldr(r5, mem[sp, 124]); // params
89 mov(r14, r3); // p = ks
90
91 // Clamp C pointers
92 cmp(r0, 2); // if mr >= 2
93 add(r4, r11, r6); // c1 = c0 + cm_stride
94 movlo(r4, r11); // c1
95 // if mr > 2
96 add(r8, r4, r6); // c2 = c1 + cm_stride
97 movls(r8, r4); // c2
98 cmp(r0, 4); // if mr >=4
99 add(r6, r8, r6); // c3 = c2 + cm_stride
100 movlo(r6, r8); // c3
101
102 // Load params values
103 vldm(mem[r5], {d12-d13}); // RNDNU params
104
105 bind(l0);
106 // Load initial bias from w into accumulators
107 vldm(mem[r9]++, {d16-d19}); // Bias
108 vmov(q10, q8);
109 vmov(q11, q9);
110 vmov(q12, q8);
111 vmov(q13, q9);
112 vmov(q14, q8);
113 vmov(q15, q9);
114
115 bind(l1);
116 // Load next 4 A pointers
117 ldr(r3, mem[r2, 0]);
118 ldr(r12, mem[r2, 4]);
119 ldr(r10, mem[r2, 8]);
120 ldr(r0, mem[r2, 12]);
121 add(r2, r2, 16);
122
123 // Add a_offset
124 ldr(r5, mem[sp, 116]); // a_offset
125 ldr(r7, mem[sp, 120]); // zero
126 cmp(r3, r7); // if a0 == zero
127 add(r3, r3, r5); // a0 += a_offset
128 moveq(r3, r7); // a0 = zero, else += a0 + a_offset
129 cmp(r12, r7); // if a1 == zero
130 add(r12, r12, r5); // a1 += a_offset
131 moveq(r12, r7); // a1 = zero, else += a1 + a_offset
132 cmp(r10, r7); // if a2 == zero
133 add(r10, r10, r5); // a2 += a_offset
134 moveq(r10, r7); // a2 = zero, else += a2 + a_offset
135 cmp(r0, r7); // if a3 == zero
136 add(r0, r0, r5); // a3 += a_offset
137 ldr(r5, mem[sp, 52]); // kc
138 moveq(r0, r7); // a3 = zero, else += a3 + a_offset
139
140 subs(r5, r5, 8); // kc - 8
141 blo(l4); // less than 8 channels?
142
143 // Main loop - 8 bytes of A.
144 // 16 SDOT, 4 LD64 A, 4 LD128 B
145 align(8);
146 bind(l2);
147 vld1_8({d0}, mem[r3]++); // A0
148 vld1_8({q2}, mem[r9]++); // B0
149 vld1_8({d1}, mem[r12]++); // A1
150 vld1_8({q3}, mem[r9]++); // B1
151 vld1_8({d2}, mem[r10]++); // A2
152 vld1_8({q4}, mem[r9]++); // B2
153 vld1_8({d3}, mem[r0]++); // A3
154 vld1_8({q5}, mem[r9]++); // B3
155 subs(r5, r5, 8);
156
157 vsdot_s8(q8, q2, d0[0]);
158 vsdot_s8(q9, q3, d0[0]);
159 vsdot_s8(q10, q2, d1[0]);
160 vsdot_s8(q11, q3, d1[0]);
161 vsdot_s8(q12, q2, d2[0]);
162 vsdot_s8(q13, q3, d2[0]);
163 vsdot_s8(q14, q2, d3[0]);
164 vsdot_s8(q15, q3, d3[0]);
165
166 vsdot_s8(q8, q4, d0[1]);
167 vsdot_s8(q9, q5, d0[1]);
168 vsdot_s8(q10, q4, d1[1]);
169 vsdot_s8(q11, q5, d1[1]);
170 vsdot_s8(q12, q4, d2[1]);
171 vsdot_s8(q13, q5, d2[1]);
172 vsdot_s8(q14, q4, d3[1]);
173 vsdot_s8(q15, q5, d3[1]);
174 bhs(l2);
175
176 // Is there a remainder?- 4 bytes of A
177 tst(r5, 4);
178 bne(l4);
179
180 bind(l3);
181 // ks loop
182 subs(r14, r14, 16); // ks -= MR * sizeof(void*)
183 bhi(l1);
184
185 ldr(r7, mem[sp, 112]); // cn_stride
186 ldr(r14, mem[sp, 56]); // p = ks
187
188 // RNDNU quantization
189 vdup_32(q0, d12[0]); // right_pre_shift
190
191 vqshl_s32(q8, q8, q0);
192 vqshl_s32(q9, q9, q0);
193 vqshl_s32(q10, q10, q0);
194 vqshl_s32(q11, q11, q0);
195 vqshl_s32(q12, q12, q0);
196 vqshl_s32(q13, q13, q0);
197 vqshl_s32(q14, q14, q0);
198 vqshl_s32(q15, q15, q0);
199
200 vdup_32(q2, d13[0]); // right_post_shift
201
202 vqdmulh_s32(q8, q8, d12[1]); // multiplier
203 vqdmulh_s32(q9, q9, d12[1]);
204 vqdmulh_s32(q10, q10, d12[1]);
205 vqdmulh_s32(q11, q11, d12[1]);
206 vqdmulh_s32(q12, q12, d12[1]);
207 vqdmulh_s32(q13, q13, d12[1]);
208 vqdmulh_s32(q14, q14, d12[1]);
209 vqdmulh_s32(q15, q15, d12[1]);
210
211 vrshl_s32(q8, q8, q2);
212 vrshl_s32(q9, q9, q2);
213 vrshl_s32(q10, q10, q2);
214 vrshl_s32(q11, q11, q2);
215 vrshl_s32(q12, q12, q2);
216 vrshl_s32(q13, q13, q2);
217 vrshl_s32(q14, q14, q2);
218 vrshl_s32(q15, q15, q2);
219 vdup_16(q0, d13[2]); // output_zero_point
220
221 vqmovn_s32(d16, q8);
222 vqmovn_s32(d17, q9);
223 vqmovn_s32(d18, q10);
224 vqmovn_s32(d19, q11);
225 vqmovn_s32(d20, q12);
226 vqmovn_s32(d21, q13);
227 vqmovn_s32(d22, q14);
228 vqmovn_s32(d23, q15);
229
230 vqadd_s16(q8, q8, q0);
231 vqadd_s16(q9, q9, q0);
232 vqadd_s16(q10, q10, q0);
233 vqadd_s16(q11, q11, q0);
234
235 vdup_8(q12, d13[6]); // output_min
236
237 vqmovn_s16(d0, q8);
238 vqmovn_s16(d1, q9);
239 vqmovn_s16(d2, q10);
240 vqmovn_s16(d3, q11);
241
242 vdup_8(q13, d13[7]); // output_max
243
244 vmax_s8(q0, q0, q12);
245 vmax_s8(q1, q1, q12);
246
247 subs(r1, r1, 8); // nc -= 8
248
249 vmin_s8(q0, q0, q13);
250 vmin_s8(q1, q1, q13);
251
252 // Store full 4 x 8
253 blo(l5);
254 vst1_8({d3}, mem[r6], r7);
255 vst1_8({d2}, mem[r8], r7);
256 vst1_8({d1}, mem[r4], r7);
257 vst1_8({d0}, mem[r11], r7);
258 sub(r2, r2, r14); // a -= ks
259 bhi(l0);
260
261 vpop({d8-d13});
262 add(sp, sp, 12); // skip pad, r2, r3
263 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
264
265 bind(l4);
266 // Remainder- 4 bytes of A
267 vld1_32({d0[0]}, mem[r3]++); // A0
268 vld1_32({q2}, mem[r9]++); // B0
269 vld1_32({d1[0]}, mem[r12]++); // A1
270 vld1_32({q3}, mem[r9]++); // B1
271 vld1_32({d2[0]}, mem[r10]++); // A2
272 vld1_32({d3[0]}, mem[r0]++); // A3
273
274 vsdot_s8(q8, q2, d0[0]);
275 vsdot_s8(q9, q3, d0[0]);
276 vsdot_s8(q10, q2, d1[0]);
277 vsdot_s8(q11, q3, d1[0]);
278 vsdot_s8(q12, q2, d2[0]);
279 vsdot_s8(q13, q3, d2[0]);
280 vsdot_s8(q14, q2, d3[0]);
281 vsdot_s8(q15, q3, d3[0]);
282 b(l3);
283
284 // Store odd width
285 align(8);
286 bind(l5);
287 tst(r1, 4);
288 beq(l6);
289 vst1_32({d3[0]}, mem[r6]++);
290 vst1_32({d2[0]}, mem[r8]++);
291 vst1_32({d1[0]}, mem[r4]++);
292 vst1_32({d0[0]}, mem[r11]++);
293 vext_8(q1, q1, q1, 4);
294 vext_8(q0, q0, q0, 4);
295 bind(l6);
296 tst(r1, 2);
297 beq(l7);
298 vst1_16({d3[0]}, mem[r6]++);
299 vst1_16({d2[0]}, mem[r8]++);
300 vst1_16({d1[0]}, mem[r4]++);
301 vst1_16({d0[0]}, mem[r11]++);
302 vext_8(q1, q1, q1, 2);
303 vext_8(q0, q0, q0, 2);
304
305 bind(l7);
306 tst(r1, 1);
307 beq(l8);
308 vst1_8({d3[0]}, mem[r6]);
309 vst1_8({d2[0]}, mem[r8]);
310 vst1_8({d1[0]}, mem[r4]);
311 vst1_8({d0[0]}, mem[r11]);
312
313 bind(l8);
314 vpop({d8-d13});
315 add(sp, sp, 12); // skip pad, r2, r3
316 pop({r4, r5, r6, r7, r8, r9, r10, r11, pc});
317 }
318 } // namespace
319 } // aarch32
320 } // xnnpack
321
xnn_generate_qs8_igemm_rndnu_ukernel_4x8c4__aarch32_neondot_ld64(xnn_code_buffer * code,size_t max_mr,size_t nc_mod_nr,size_t kc,size_t ks,const void * params)322 xnn_status_t xnn_generate_qs8_igemm_rndnu_ukernel_4x8c4__aarch32_neondot_ld64(xnn_code_buffer* code, size_t max_mr, size_t nc_mod_nr, size_t kc, size_t ks, const void* params) {
323 using namespace xnnpack::aarch32;
324 Generator g(code);
325 g.generate(max_mr, nc_mod_nr, kc, ks, nullptr);
326 g.finalize();
327 if (g.error() != xnnpack::Error::kNoError) {
328 return xnn_status_invalid_state;
329 }
330 return xnn_status_success;
331 }
332