1 //
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 // http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13
14 // fixedpoint_avx.h: optimized avx specializations of the templates
15 // in fixedpoint.h.
16
17 #ifndef GEMMLOWP_INTERNAL_FIXEDPOINT_AVX_H_
18 #define GEMMLOWP_INTERNAL_FIXEDPOINT_AVX_H_
19
20 #include <immintrin.h>
21 #include "fixedpoint.h"
22 #include "fixedpoint_sse.h"
23
24 namespace gemmlowp {
25
26 struct int16x16_m256i {
27 __m256i v;
28 };
29
30 // Keep int16x16_m256i trivially constructible/destructible and provide
31 // easily optimized helper function.
to_int16x16_m256i(__m256i w)32 inline int16x16_m256i to_int16x16_m256i(__m256i w) {
33 int16x16_m256i r;
34 r.v = w;
35 return r;
36 }
37
38 template <>
39 struct FixedPointRawTypeTraits<__m256i> {
40 typedef std::int32_t ScalarRawType;
41 // TODO: This can actually support up to 8 lanes, so we should either
42 // change to 8 or create int32x8_m256i struct to handle that case.
43 static const int kLanes = 4;
44 };
45
46 template <>
47 struct FixedPointRawTypeTraits<int16x16_m256i> {
48 typedef std::int16_t ScalarRawType;
49 static const int kLanes = 16;
50 };
51
52 template <>
53 inline __m256i BitAnd(__m256i a, __m256i b) {
54 return _mm256_and_si256(a, b);
55 }
56
57 template <>
58 inline int16x16_m256i BitAnd(int16x16_m256i a, int16x16_m256i b) {
59 return to_int16x16_m256i(_mm256_and_si256(a.v, b.v));
60 }
61
62 template <>
63 inline __m256i BitOr(__m256i a, __m256i b) {
64 return _mm256_or_si256(a, b);
65 }
66
67 template <>
68 inline int16x16_m256i BitOr(int16x16_m256i a, int16x16_m256i b) {
69 return to_int16x16_m256i(_mm256_or_si256(a.v, b.v));
70 }
71
72 template <>
73 inline __m256i BitXor(__m256i a, __m256i b) {
74 return _mm256_xor_si256(a, b);
75 }
76
77 template <>
78 inline int16x16_m256i BitXor(int16x16_m256i a, int16x16_m256i b) {
79 return to_int16x16_m256i(_mm256_xor_si256(a.v, b.v));
80 }
81
82 template <>
83 inline __m256i BitNot(__m256i a) {
84 return _mm256_andnot_si256(a, _mm256_set1_epi32(-1));
85 }
86
87 template <>
88 inline int16x16_m256i BitNot(int16x16_m256i a) {
89 return to_int16x16_m256i(_mm256_andnot_si256(a.v, _mm256_set1_epi16(-1)));
90 }
91
92 template <>
93 inline __m256i Add(__m256i a, __m256i b) {
94 return _mm256_add_epi32(a, b);
95 }
96
97 template <>
98 inline int16x16_m256i Add(int16x16_m256i a, int16x16_m256i b) {
99 return to_int16x16_m256i(_mm256_add_epi16(a.v, b.v));
100 }
101
102 template <>
103 inline __m256i Mul(__m256i a, __m256i b) {
104 return _mm256_mullo_epi32(a, b);
105 }
106
107 template <>
108 inline int16x16_m256i Mul(int16x16_m256i a, int16x16_m256i b) {
109 return to_int16x16_m256i(_mm256_mullo_epi16(a.v, b.v));
110 }
111
112 template <>
113 inline __m256i Sub(__m256i a, __m256i b) {
114 return _mm256_sub_epi32(a, b);
115 }
116
117 template <>
118 inline int16x16_m256i Sub(int16x16_m256i a, int16x16_m256i b) {
119 return to_int16x16_m256i(_mm256_sub_epi16(a.v, b.v));
120 }
121
122 template <>
123 inline __m256i Neg(__m256i a) {
124 return _mm256_sign_epi32(a, _mm256_set1_epi32(-1));
125 }
126
127 template <>
128 inline int16x16_m256i Neg(int16x16_m256i a) {
129 return to_int16x16_m256i(_mm256_sign_epi16(a.v, _mm256_set1_epi16(-1)));
130 }
131
132 template <>
133 inline __m256i ShiftLeft(__m256i a, int offset) {
134 return _mm256_slli_epi32(a, offset);
135 }
136
137 template <>
138 inline int16x16_m256i ShiftLeft(int16x16_m256i a, int offset) {
139 return to_int16x16_m256i(_mm256_slli_epi16(a.v, offset));
140 }
141
142 template <>
143 inline __m256i ShiftRight(__m256i a, int offset) {
144 return _mm256_srai_epi32(a, offset);
145 }
146
147 template <>
148 inline int16x16_m256i ShiftRight(int16x16_m256i a, int offset) {
149 return to_int16x16_m256i(_mm256_srai_epi16(a.v, offset));
150 }
151
152 template <>
153 inline __m256i SelectUsingMask(__m256i if_mask, __m256i then_val,
154 __m256i else_val) {
155 return _mm256_castps_si256(_mm256_blendv_ps(_mm256_castsi256_ps(else_val),
156 _mm256_castsi256_ps(then_val),
157 _mm256_castsi256_ps(if_mask)));
158 }
159
160 template <>
161 inline int16x16_m256i SelectUsingMask(int16x16_m256i if_mask,
162 int16x16_m256i then_val,
163 int16x16_m256i else_val) {
164 // Borrowed from Intel's arm_neon_sse.h header.
165 return to_int16x16_m256i(
166 _mm256_or_si256(_mm256_and_si256(if_mask.v, then_val.v),
167 _mm256_andnot_si256(if_mask.v, else_val.v)));
168 }
169
170 template <>
171 inline __m256i MaskIfEqual(__m256i a, __m256i b) {
172 return _mm256_cmpeq_epi32(a, b);
173 }
174
175 template <>
176 inline int16x16_m256i MaskIfEqual(int16x16_m256i a, int16x16_m256i b) {
177 return to_int16x16_m256i(_mm256_cmpeq_epi16(a.v, b.v));
178 }
179
180 template <>
181 inline __m256i MaskIfNotEqual(__m256i a, __m256i b) {
182 return BitNot(MaskIfEqual(a, b));
183 }
184
185 template <>
186 inline int16x16_m256i MaskIfNotEqual(int16x16_m256i a, int16x16_m256i b) {
187 return BitNot(MaskIfEqual(a, b));
188 }
189
190 template <>
191 inline __m256i MaskIfZero(__m256i a) {
192 return MaskIfEqual(a, _mm256_set1_epi32(0));
193 }
194
195 template <>
196 inline int16x16_m256i MaskIfZero(int16x16_m256i a) {
197 return MaskIfEqual(a, to_int16x16_m256i(_mm256_set1_epi16(0)));
198 }
199
200 template <>
201 inline __m256i MaskIfNonZero(__m256i a) {
202 return MaskIfNotEqual(a, _mm256_set1_epi32(0));
203 }
204
205 template <>
206 inline int16x16_m256i MaskIfNonZero(int16x16_m256i a) {
207 return MaskIfNotEqual(a, to_int16x16_m256i(_mm256_set1_epi16(0)));
208 }
209
210 template <>
211 inline __m256i MaskIfGreaterThan(__m256i a, __m256i b) {
212 return _mm256_cmpgt_epi32(a, b);
213 }
214
215 template <>
216 inline int16x16_m256i MaskIfGreaterThan(int16x16_m256i a, int16x16_m256i b) {
217 return to_int16x16_m256i(_mm256_cmpgt_epi16(a.v, b.v));
218 }
219
220 template <>
221 inline __m256i MaskIfLessThan(__m256i a, __m256i b) {
222 return _mm256_cmpgt_epi32(b, a);
223 }
224
225 template <>
226 inline int16x16_m256i MaskIfLessThan(int16x16_m256i a, int16x16_m256i b) {
227 return to_int16x16_m256i(_mm256_cmpgt_epi16(b.v, a.v));
228 }
229
230 template <>
231 inline __m256i MaskIfGreaterThanOrEqual(__m256i a, __m256i b) {
232 return BitNot(MaskIfLessThan(a, b));
233 }
234
235 template <>
236 inline int16x16_m256i MaskIfGreaterThanOrEqual(int16x16_m256i a,
237 int16x16_m256i b) {
238 return BitNot(MaskIfLessThan(a, b));
239 }
240
241 template <>
242 inline __m256i MaskIfLessThanOrEqual(__m256i a, __m256i b) {
243 return BitNot(MaskIfGreaterThan(a, b));
244 }
245
246 template <>
247 inline int16x16_m256i MaskIfLessThanOrEqual(int16x16_m256i a,
248 int16x16_m256i b) {
249 return BitNot(MaskIfGreaterThan(a, b));
250 }
251
252 /* Assumptions:
253 - All and Any are used on masks.
254 - masks are all_ones for true lanes, all_zeroes otherwise.
255 Hence, All means all 128bits set, and Any means any bit set.
256 */
257
258 template <>
259 inline bool All(__m256i a) {
260 return _mm256_testc_si256(a, a);
261 }
262
263 template <>
264 inline bool All(int16x16_m256i a) {
265 return _mm256_testc_si256(a.v, a.v);
266 }
267
268 template <>
269 inline bool Any(__m256i a) {
270 return BitNot(_mm256_testz_si256(a, a));
271 }
272
273 template <>
274 inline bool Any(int16x16_m256i a) {
275 return BitNot(_mm256_testz_si256(a.v, a.v));
276 }
277
278 template <>
279 inline __m256i RoundingHalfSum(__m256i a, __m256i b) {
280 /* __m256i round_bit_mask, a_over_2, b_over_2, round_bit, sum; */
281 /* We divide the inputs before the add to avoid the overflow and costly test
282 */
283 /* of checking if an overflow occured on signed add */
284 /* round_bit_mask = _mm_set1_epi32(1); */
285 /* a_over_2 = _mm_srai_epi32(a, 1); */
286 /* b_over_2 = _mm_srai_epi32(b, 1); */
287 /* sum = Add(a_over_2, b_over_2); */
288 /* round_bit = _mm_sign_epi32(BitAnd(BitOr(a,b), round_bit_mask), sum); */
289 /* return Add(sum, round_bit); */
290
291 /* Other possibility detecting overflow and xor the sign if an overflow
292 * happened*/
293 __m256i one, sign_bit_mask, sum, rounded_half_sum, overflow, result;
294 one = _mm256_set1_epi32(1);
295 sign_bit_mask = _mm256_set1_epi32(0x80000000);
296 sum = Add(a, b);
297 rounded_half_sum = _mm256_srai_epi32(Add(sum, one), 1);
298 overflow =
299 BitAnd(BitAnd(BitXor(a, rounded_half_sum), BitXor(b, rounded_half_sum)),
300 sign_bit_mask);
301 result = BitXor(rounded_half_sum, overflow);
302 return result;
303 }
304
305 template <>
306 inline int16x16_m256i RoundingHalfSum(int16x16_m256i a, int16x16_m256i b) {
307 // Borrowed from Intel's arm_neon_sse.h header.
308 __m256i constant_neg_32768 = _mm256_set1_epi16(-32768);
309 __m256i a_unsigned = _mm256_sub_epi16(a.v, constant_neg_32768);
310 __m256i b_unsigned = _mm256_sub_epi16(b.v, constant_neg_32768);
311 __m256i avg_unsigned = _mm256_avg_epu16(a_unsigned, b_unsigned);
312 __m256i avg = _mm256_add_epi16(avg_unsigned, constant_neg_32768);
313 return to_int16x16_m256i(avg);
314 }
315
316 template <>
317 inline __m256i SaturatingRoundingDoublingHighMul(__m256i a, __m256i b) {
318 __m256i min, saturation_mask, a0_a2, a1_a3, b0_b2, b1_b3;
319 __m256i a0b0_a2b2, a1b1_a3b3, a0b0_a2b2_rounded, a1b1_a3b3_rounded;
320 __m256i a0b0_a2b2_rounded_2x, a1b1_a3b3_rounded_2x, result;
321 __m256i nudge;
322
323 // saturation only happen if a == b == INT_MIN
324 min = _mm256_set1_epi32(std::numeric_limits<std::int32_t>::min());
325 saturation_mask = BitAnd(MaskIfEqual(a, b), MaskIfEqual(a, min));
326
327 // a = a0 | a1 | a2 | a3
328 // b = b0 | b1 | b2 | b3
329 a0_a2 = a;
330 a1_a3 = _mm256_srli_si256(a, 4);
331 b0_b2 = b;
332 b1_b3 = _mm256_srli_si256(b, 4);
333
334 a0b0_a2b2 = _mm256_mul_epi32(a0_a2, b0_b2);
335 a1b1_a3b3 = _mm256_mul_epi32(a1_a3, b1_b3);
336
337 // do the rounding and take into account that it will be doubled
338 nudge = _mm256_set1_epi64x(1 << 30);
339 a0b0_a2b2_rounded = _mm256_add_epi64(a0b0_a2b2, nudge);
340 a1b1_a3b3_rounded = _mm256_add_epi64(a1b1_a3b3, nudge);
341
342 // do the doubling
343 a0b0_a2b2_rounded_2x = _mm256_slli_epi64(a0b0_a2b2_rounded, 1);
344 a1b1_a3b3_rounded_2x = _mm256_slli_epi64(a1b1_a3b3_rounded, 1);
345
346 // get the high part of the products
347 result = _mm256_blend_epi16(_mm256_srli_si256(a0b0_a2b2_rounded_2x, 4),
348 a1b1_a3b3_rounded_2x, 0xcc);
349
350 // saturate those which overflowed
351 return SelectUsingMask(saturation_mask, min, result);
352 }
353
354 template <>
355 inline int16x16_m256i SaturatingRoundingDoublingHighMul(int16x16_m256i a,
356 int16x16_m256i b) {
357 // Use _mm256_mulhrs_epi16 then saturate with a bit-operation,
358 // borrowed from Intel's arm_neon_sse.h header.
359 __m256i result_unsaturated = _mm256_mulhrs_epi16(a.v, b.v);
360 __m256i saturation_mask =
361 _mm256_cmpeq_epi16(result_unsaturated, _mm256_set1_epi16(0x8000));
362 __m256i result = _mm256_xor_si256(result_unsaturated, saturation_mask);
363 return to_int16x16_m256i(result);
364 }
365
366 template <>
367 inline __m256i Dup<__m256i>(std::int32_t x) {
368 return _mm256_set1_epi32(x);
369 }
370
371 template <>
372 inline int16x16_m256i Dup<int16x16_m256i>(std::int16_t x) {
373 return to_int16x16_m256i(_mm256_set1_epi16(x));
374 }
375
376 // So far this is only needed for int16.
377 template <>
378 inline int16x16_m256i SaturatingAdd(int16x16_m256i a, int16x16_m256i b) {
379 return to_int16x16_m256i(_mm256_adds_epi16(a.v, b.v));
380 }
381
382 } // end namespace gemmlowp
383
384 #endif // GEMMLOWP_INTERNAL_FIXEDPOINT_AVX_H_
385