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