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