RapidFuzz
Loading...
Searching...
No Matches
simd_avx2.hpp
1/* SPDX-License-Identifier: MIT */
2/* Copyright © 2022 Max Bachmann */
3#pragma once
4
5#include <array>
6#include <immintrin.h>
7#include <ostream>
8#include <rapidfuzz/details/intrinsics.hpp>
9#include <stdint.h>
10
11namespace rapidfuzz {
12namespace detail {
13namespace simd_avx2 {
14
15template <typename T>
16class native_simd;
17
18template <>
19class native_simd<uint64_t> {
20public:
21 using value_type = uint64_t;
22
23 static constexpr int alignment = 32;
24 static const int size = 4;
25 __m256i xmm;
26
27 native_simd() noexcept
28 {}
29
30 native_simd(__m256i val) noexcept : xmm(val)
31 {}
32
33 native_simd(uint64_t a) noexcept
34 {
35 xmm = _mm256_set1_epi64x(static_cast<int64_t>(a));
36 }
37
38 native_simd(const uint64_t* p) noexcept
39 {
40 load(p);
41 }
42
43 operator __m256i() const noexcept
44 {
45 return xmm;
46 }
47
48 native_simd load(const uint64_t* p) noexcept
49 {
50 xmm = _mm256_set_epi64x(static_cast<int64_t>(p[3]), static_cast<int64_t>(p[2]),
51 static_cast<int64_t>(p[1]), static_cast<int64_t>(p[0]));
52 return *this;
53 }
54
55 void store(uint64_t* p) const noexcept
56 {
57 _mm256_store_si256(reinterpret_cast<__m256i*>(p), xmm);
58 }
59
60 native_simd operator+(const native_simd b) const noexcept
61 {
62 return _mm256_add_epi64(xmm, b);
63 }
64
65 native_simd& operator+=(const native_simd b) noexcept
66 {
67 xmm = _mm256_add_epi64(xmm, b);
68 return *this;
69 }
70
71 native_simd operator-(const native_simd b) const noexcept
72 {
73 return _mm256_sub_epi64(xmm, b);
74 }
75
76 native_simd operator-() const noexcept
77 {
78 return _mm256_sub_epi64(_mm256_setzero_si256(), xmm);
79 }
80
81 native_simd& operator-=(const native_simd b) noexcept
82 {
83 xmm = _mm256_sub_epi64(xmm, b);
84 return *this;
85 }
86};
87
88template <>
89class native_simd<uint32_t> {
90public:
91 using value_type = uint32_t;
92
93 static constexpr int alignment = 32;
94 static const int size = 8;
95 __m256i xmm;
96
97 native_simd() noexcept
98 {}
99
100 native_simd(__m256i val) noexcept : xmm(val)
101 {}
102
103 native_simd(uint32_t a) noexcept
104 {
105 xmm = _mm256_set1_epi32(static_cast<int>(a));
106 }
107
108 native_simd(const uint64_t* p) noexcept
109 {
110 load(p);
111 }
112
113 operator __m256i() const
114 {
115 return xmm;
116 }
117
118 native_simd load(const uint64_t* p) noexcept
119 {
120 xmm = _mm256_set_epi64x(static_cast<int64_t>(p[3]), static_cast<int64_t>(p[2]),
121 static_cast<int64_t>(p[1]), static_cast<int64_t>(p[0]));
122 return *this;
123 }
124
125 void store(uint32_t* p) const noexcept
126 {
127 _mm256_store_si256(reinterpret_cast<__m256i*>(p), xmm);
128 }
129
130 native_simd operator+(const native_simd b) const noexcept
131 {
132 return _mm256_add_epi32(xmm, b);
133 }
134
135 native_simd& operator+=(const native_simd b) noexcept
136 {
137 xmm = _mm256_add_epi32(xmm, b);
138 return *this;
139 }
140
141 native_simd operator-() const noexcept
142 {
143 return _mm256_sub_epi32(_mm256_setzero_si256(), xmm);
144 }
145
146 native_simd operator-(const native_simd b) const noexcept
147 {
148 return _mm256_sub_epi32(xmm, b);
149 }
150
151 native_simd& operator-=(const native_simd b) noexcept
152 {
153 xmm = _mm256_sub_epi32(xmm, b);
154 return *this;
155 }
156};
157
158template <>
159class native_simd<uint16_t> {
160public:
161 using value_type = uint16_t;
162
163 static constexpr int alignment = 32;
164 static const int size = 16;
165 __m256i xmm;
166
167 native_simd() noexcept
168 {}
169
170 native_simd(__m256i val) : xmm(val)
171 {}
172
173 native_simd(uint16_t a) noexcept
174 {
175 xmm = _mm256_set1_epi16(static_cast<short>(a));
176 }
177
178 native_simd(const uint64_t* p) noexcept
179 {
180 load(p);
181 }
182
183 operator __m256i() const noexcept
184 {
185 return xmm;
186 }
187
188 native_simd load(const uint64_t* p) noexcept
189 {
190 xmm = _mm256_set_epi64x(static_cast<int64_t>(p[3]), static_cast<int64_t>(p[2]),
191 static_cast<int64_t>(p[1]), static_cast<int64_t>(p[0]));
192 return *this;
193 }
194
195 void store(uint16_t* p) const noexcept
196 {
197 _mm256_store_si256(reinterpret_cast<__m256i*>(p), xmm);
198 }
199
200 native_simd operator+(const native_simd b) const noexcept
201 {
202 return _mm256_add_epi16(xmm, b);
203 }
204
205 native_simd& operator+=(const native_simd b) noexcept
206 {
207 xmm = _mm256_add_epi16(xmm, b);
208 return *this;
209 }
210
211 native_simd operator-(const native_simd b) const noexcept
212 {
213 return _mm256_sub_epi16(xmm, b);
214 }
215
216 native_simd operator-() const noexcept
217 {
218 return _mm256_sub_epi16(_mm256_setzero_si256(), xmm);
219 }
220
221 native_simd& operator-=(const native_simd b) noexcept
222 {
223 xmm = _mm256_sub_epi16(xmm, b);
224 return *this;
225 }
226};
227
228template <>
229class native_simd<uint8_t> {
230public:
231 using value_type = uint8_t;
232
233 static constexpr int alignment = 32;
234 static const int size = 32;
235 __m256i xmm;
236
237 native_simd() noexcept
238 {}
239
240 native_simd(__m256i val) noexcept : xmm(val)
241 {}
242
243 native_simd(uint8_t a) noexcept
244 {
245 xmm = _mm256_set1_epi8(static_cast<char>(a));
246 }
247
248 native_simd(const uint64_t* p) noexcept
249 {
250 load(p);
251 }
252
253 operator __m256i() const noexcept
254 {
255 return xmm;
256 }
257
258 native_simd load(const uint64_t* p) noexcept
259 {
260 xmm = _mm256_set_epi64x(static_cast<int64_t>(p[3]), static_cast<int64_t>(p[2]),
261 static_cast<int64_t>(p[1]), static_cast<int64_t>(p[0]));
262 return *this;
263 }
264
265 void store(uint8_t* p) const noexcept
266 {
267 _mm256_store_si256(reinterpret_cast<__m256i*>(p), xmm);
268 }
269
270 native_simd operator+(const native_simd b) const noexcept
271 {
272 return _mm256_add_epi8(xmm, b);
273 }
274
275 native_simd& operator+=(const native_simd b) noexcept
276 {
277 xmm = _mm256_add_epi8(xmm, b);
278 return *this;
279 }
280
281 native_simd operator-(const native_simd b) const noexcept
282 {
283 return _mm256_sub_epi8(xmm, b);
284 }
285
286 native_simd operator-() const noexcept
287 {
288 return _mm256_sub_epi8(_mm256_setzero_si256(), xmm);
289 }
290
291 native_simd& operator-=(const native_simd b) noexcept
292 {
293 xmm = _mm256_sub_epi8(xmm, b);
294 return *this;
295 }
296};
297
298template <typename T>
299std::ostream& operator<<(std::ostream& os, const native_simd<T>& a)
300{
301 alignas(native_simd<T>::alignment) std::array<T, native_simd<T>::size> res;
302 a.store(&res[0]);
303
304 for (size_t i = res.size() - 1; i != 0; i--)
305 os << std::bitset<std::numeric_limits<T>::digits>(res[i]) << "|";
306
307 os << std::bitset<std::numeric_limits<T>::digits>(res[0]);
308 return os;
309}
310
311template <typename T>
312__m256i hadd_impl(__m256i x) noexcept;
313
314template <>
315inline __m256i hadd_impl<uint8_t>(__m256i x) noexcept
316{
317 return x;
318}
319
320template <>
321inline __m256i hadd_impl<uint16_t>(__m256i x) noexcept
322{
323 const __m256i mask = _mm256_set1_epi16(0x001f);
324 __m256i y = _mm256_srli_si256(x, 1);
325 x = _mm256_add_epi16(x, y);
326 return _mm256_and_si256(x, mask);
327}
328
329template <>
330inline __m256i hadd_impl<uint32_t>(__m256i x) noexcept
331{
332 const __m256i mask = _mm256_set1_epi32(0x0000003F);
333 x = hadd_impl<uint16_t>(x);
334 __m256i y = _mm256_srli_si256(x, 2);
335 x = _mm256_add_epi32(x, y);
336 return _mm256_and_si256(x, mask);
337}
338
339template <>
340inline __m256i hadd_impl<uint64_t>(__m256i x) noexcept
341{
342 return _mm256_sad_epu8(x, _mm256_setzero_si256());
343}
344
345/* based on the paper `Faster Population Counts Using AVX2 Instructions` */
346template <typename T>
347native_simd<T> popcount_impl(const native_simd<T>& v) noexcept
348{
349 __m256i lookup = _mm256_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3,
350 1, 2, 2, 3, 2, 3, 3, 4);
351 const __m256i low_mask = _mm256_set1_epi8(0x0F);
352 __m256i lo = _mm256_and_si256(v, low_mask);
353 __m256i hi = _mm256_and_si256(_mm256_srli_epi32(v, 4), low_mask);
354 __m256i popcnt1 = _mm256_shuffle_epi8(lookup, lo);
355 __m256i popcnt2 = _mm256_shuffle_epi8(lookup, hi);
356 __m256i total = _mm256_add_epi8(popcnt1, popcnt2);
357 return hadd_impl<T>(total);
358}
359
360template <typename T>
361std::array<T, native_simd<T>::size> popcount(const native_simd<T>& a) noexcept
362{
363 alignas(native_simd<T>::alignment) std::array<T, native_simd<T>::size> res;
364 popcount_impl(a).store(&res[0]);
365 return res;
366}
367
368// function andnot: a & ~ b
369template <typename T>
370native_simd<T> andnot(const native_simd<T>& a, const native_simd<T>& b)
371{
372 return _mm256_andnot_si256(b, a);
373}
374
375static inline native_simd<uint8_t> operator==(const native_simd<uint8_t>& a,
376 const native_simd<uint8_t>& b) noexcept
377{
378 return _mm256_cmpeq_epi8(a, b);
379}
380
381static inline native_simd<uint16_t> operator==(const native_simd<uint16_t>& a,
382 const native_simd<uint16_t>& b) noexcept
383{
384 return _mm256_cmpeq_epi16(a, b);
385}
386
387static inline native_simd<uint32_t> operator==(const native_simd<uint32_t>& a,
388 const native_simd<uint32_t>& b) noexcept
389{
390 return _mm256_cmpeq_epi32(a, b);
391}
392
393static inline native_simd<uint64_t> operator==(const native_simd<uint64_t>& a,
394 const native_simd<uint64_t>& b) noexcept
395{
396 return _mm256_cmpeq_epi64(a, b);
397}
398
399template <typename T>
400static inline native_simd<T> operator!=(const native_simd<T>& a, const native_simd<T>& b) noexcept
401{
402 return ~(a == b);
403}
404
405static inline native_simd<uint8_t> operator<<(const native_simd<uint8_t>& a, int b) noexcept
406{
407 char mask = static_cast<char>(0xFF >> b);
408 __m256i am = _mm256_and_si256(a, _mm256_set1_epi8(mask));
409 return _mm256_slli_epi16(am, b);
410}
411
412static inline native_simd<uint16_t> operator<<(const native_simd<uint16_t>& a, int b) noexcept
413{
414 return _mm256_slli_epi16(a, b);
415}
416
417static inline native_simd<uint32_t> operator<<(const native_simd<uint32_t>& a, int b) noexcept
418{
419 return _mm256_slli_epi32(a, b);
420}
421
422static inline native_simd<uint64_t> operator<<(const native_simd<uint64_t>& a, int b) noexcept
423{
424 return _mm256_slli_epi64(a, b);
425}
426
427static inline native_simd<uint8_t> operator>>(const native_simd<uint8_t>& a, int b) noexcept
428{
429 char mask = static_cast<char>(0xFF << b);
430 __m256i am = _mm256_and_si256(a, _mm256_set1_epi8(mask));
431 return _mm256_srli_epi16(am, b);
432}
433
434static inline native_simd<uint16_t> operator>>(const native_simd<uint16_t>& a, int b) noexcept
435{
436 return _mm256_srli_epi16(a, b);
437}
438
439static inline native_simd<uint32_t> operator>>(const native_simd<uint32_t>& a, int b) noexcept
440{
441 return _mm256_srli_epi32(a, b);
442}
443
444static inline native_simd<uint64_t> operator>>(const native_simd<uint64_t>& a, int b) noexcept
445{
446 return _mm256_srli_epi64(a, b);
447}
448
449template <typename T>
450native_simd<T> operator&(const native_simd<T>& a, const native_simd<T>& b) noexcept
451{
452 return _mm256_and_si256(a, b);
453}
454
455template <typename T>
456native_simd<T> operator&=(native_simd<T>& a, const native_simd<T>& b) noexcept
457{
458 a = a & b;
459 return a;
460}
461
462template <typename T>
463native_simd<T> operator|(const native_simd<T>& a, const native_simd<T>& b) noexcept
464{
465 return _mm256_or_si256(a, b);
466}
467
468template <typename T>
469native_simd<T> operator|=(native_simd<T>& a, const native_simd<T>& b) noexcept
470{
471 a = a | b;
472 return a;
473}
474
475template <typename T>
476native_simd<T> operator^(const native_simd<T>& a, const native_simd<T>& b) noexcept
477{
478 return _mm256_xor_si256(a, b);
479}
480
481template <typename T>
482native_simd<T> operator^=(native_simd<T>& a, const native_simd<T>& b) noexcept
483{
484 a = a ^ b;
485 return a;
486}
487
488template <typename T>
489native_simd<T> operator~(const native_simd<T>& a) noexcept
490{
491 return _mm256_xor_si256(a, _mm256_set1_epi32(-1));
492}
493
494// potentially we want a special native_simd<bool> for this
495static inline native_simd<uint8_t> operator>=(const native_simd<uint8_t>& a,
496 const native_simd<uint8_t>& b) noexcept
497{
498 return _mm256_cmpeq_epi8(_mm256_max_epu8(a, b), a); // a == max(a,b)
499}
500
501static inline native_simd<uint16_t> operator>=(const native_simd<uint16_t>& a,
502 const native_simd<uint16_t>& b) noexcept
503{
504 return _mm256_cmpeq_epi16(_mm256_max_epu16(a, b), a); // a == max(a,b)
505}
506
507static inline native_simd<uint32_t> operator>=(const native_simd<uint32_t>& a,
508 const native_simd<uint32_t>& b) noexcept
509{
510 return _mm256_cmpeq_epi32(_mm256_max_epu32(a, b), a); // a == max(a,b)
511}
512
513static inline native_simd<uint64_t> operator>(const native_simd<uint64_t>& a,
514 const native_simd<uint64_t>& b) noexcept;
515
516static inline native_simd<uint64_t> operator>=(const native_simd<uint64_t>& a,
517 const native_simd<uint64_t>& b) noexcept
518{
519 return ~(b > a);
520}
521
522template <typename T>
523static inline native_simd<T> operator<=(const native_simd<T>& a, const native_simd<T>& b) noexcept
524{
525 return b >= a;
526}
527
528static inline native_simd<uint8_t> operator>(const native_simd<uint8_t>& a,
529 const native_simd<uint8_t>& b) noexcept
530{
531 return ~(b >= a);
532}
533
534static inline native_simd<uint16_t> operator>(const native_simd<uint16_t>& a,
535 const native_simd<uint16_t>& b) noexcept
536{
537 return ~(b >= a);
538}
539
540static inline native_simd<uint32_t> operator>(const native_simd<uint32_t>& a,
541 const native_simd<uint32_t>& b) noexcept
542{
543 __m256i signbit = _mm256_set1_epi32(static_cast<int32_t>(0x80000000));
544 __m256i a1 = _mm256_xor_si256(a, signbit);
545 __m256i b1 = _mm256_xor_si256(b, signbit);
546 return _mm256_cmpgt_epi32(a1, b1); // signed compare
547}
548
549static inline native_simd<uint64_t> operator>(const native_simd<uint64_t>& a,
550 const native_simd<uint64_t>& b) noexcept
551{
552 __m256i sign64 = native_simd<uint64_t>(0x8000000000000000);
553 __m256i aflip = _mm256_xor_si256(a, sign64);
554 __m256i bflip = _mm256_xor_si256(b, sign64);
555 return _mm256_cmpgt_epi64(aflip, bflip); // signed compare
556}
557
558template <typename T>
559static inline native_simd<T> operator<(const native_simd<T>& a, const native_simd<T>& b) noexcept
560{
561 return b > a;
562}
563
564template <typename T>
565static inline native_simd<T> max8(const native_simd<T>& a, const native_simd<T>& b) noexcept
566{
567 return _mm256_max_epu8(a, b);
568}
569
570template <typename T>
571static inline native_simd<T> max16(const native_simd<T>& a, const native_simd<T>& b) noexcept
572{
573 return _mm256_max_epu16(a, b);
574}
575
576template <typename T>
577static inline native_simd<T> max32(const native_simd<T>& a, const native_simd<T>& b) noexcept
578{
579 return _mm256_max_epu32(a, b);
580}
581
582template <typename T>
583static inline native_simd<T> min8(const native_simd<T>& a, const native_simd<T>& b) noexcept
584{
585 return _mm256_min_epu8(a, b);
586}
587
588template <typename T>
589static inline native_simd<T> min16(const native_simd<T>& a, const native_simd<T>& b) noexcept
590{
591 return _mm256_min_epu16(a, b);
592}
593
594template <typename T>
595static inline native_simd<T> min32(const native_simd<T>& a, const native_simd<T>& b) noexcept
596{
597 return _mm256_min_epu32(a, b);
598}
599
600/* taken from https://stackoverflow.com/a/51807800 */
601static inline native_simd<uint8_t> sllv(const native_simd<uint8_t>& a,
602 const native_simd<uint8_t>& count_) noexcept
603{
604 __m256i mask_hi = _mm256_set1_epi32(static_cast<int32_t>(0xFF00FF00));
605 __m256i multiplier_lut = _mm256_set_epi8(0, 0, 0, 0, 0, 0, 0, 0, char(-128), 64, 32, 16, 8, 4, 2, 1, 0, 0,
606 0, 0, 0, 0, 0, 0, char(-128), 64, 32, 16, 8, 4, 2, 1);
607
608 __m256i count_sat =
609 _mm256_min_epu8(count_, _mm256_set1_epi8(8)); /* AVX shift counts are not masked. So a_i << n_i = 0
610 for n_i >= 8. count_sat is always less than 9.*/
611 __m256i multiplier = _mm256_shuffle_epi8(
612 multiplier_lut, count_sat); /* Select the right multiplication factor in the lookup table. */
613 __m256i x_lo = _mm256_mullo_epi16(a, multiplier); /* Unfortunately _mm256_mullo_epi8 doesn't exist. Split
614 the 16 bit elements in a high and low part. */
615
616 __m256i multiplier_hi = _mm256_srli_epi16(multiplier, 8); /* The multiplier of the high bits. */
617 __m256i a_hi = _mm256_and_si256(a, mask_hi); /* Mask off the low bits. */
618 __m256i x_hi = _mm256_mullo_epi16(a_hi, multiplier_hi);
619 __m256i x = _mm256_blendv_epi8(x_lo, x_hi, mask_hi); /* Merge the high and low part. */
620 return x;
621}
622
623/* taken from https://stackoverflow.com/a/51805592 */
624static inline native_simd<uint16_t> sllv(const native_simd<uint16_t>& a,
625 const native_simd<uint16_t>& count) noexcept
626{
627 const __m256i mask = _mm256_set1_epi32(static_cast<int32_t>(0xFFFF0000));
628 __m256i low_half = _mm256_sllv_epi32(a, _mm256_andnot_si256(mask, count));
629 __m256i high_half = _mm256_sllv_epi32(_mm256_and_si256(mask, a), _mm256_srli_epi32(count, 16));
630 return _mm256_blend_epi16(low_half, high_half, 0xAA);
631}
632
633static inline native_simd<uint32_t> sllv(const native_simd<uint32_t>& a,
634 const native_simd<uint32_t>& count) noexcept
635{
636 return _mm256_sllv_epi32(a, count);
637}
638
639static inline native_simd<uint64_t> sllv(const native_simd<uint64_t>& a,
640 const native_simd<uint64_t>& count) noexcept
641{
642 return _mm256_sllv_epi64(a, count);
643}
644
645} // namespace simd_avx2
646} // namespace detail
647} // namespace rapidfuzz
Generated by doxygen 1.9.8