Skip to content

Commit 1134be2

Browse files
rafaelcfsousarafaelcfsousa
committed
SIMD: Use universal intrinsics to implement comparison functions
1 parent ae8b9ce commit 1134be2

18 files changed

Lines changed: 852 additions & 200 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,7 @@ numpy/core/src/umath/loops_exponent_log.dispatch.c
225225
numpy/core/src/umath/loops_umath_fp.dispatch.c
226226
numpy/core/src/umath/loops_hyperbolic.dispatch.c
227227
numpy/core/src/umath/loops_modulo.dispatch.c
228+
numpy/core/src/umath/loops_comparison.dispatch.c
228229
# npysort module
229230
numpy/core/src/npysort/x86-qsort.dispatch.c
230231
numpy/core/src/npysort/x86-qsort.dispatch.*.cpp

‎benchmarks/benchmarks/bench_ufunc.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,8 +170,25 @@ def time_divide_scalar2(self, dtype):
170170
def time_divide_scalar2_inplace(self, dtype):
171171
np.divide(self.d, 1, out=self.d)
172172

173+
174+
class CustomComparison(Benchmark):
175+
params = (np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16,
176+
np.uint32, np.uint64, np.float32, np.float64, np.bool_)
177+
param_names = ['dtype']
178+
179+
def setup(self, dtype):
180+
self.x = np.ones(50000, dtype=dtype)
181+
self.y = np.ones(50000, dtype=dtype)
182+
self.s = np.ones(1, dtype=dtype)
183+
184+
def time_less_than_binary(self, dtype):
185+
(self.x < self.y)
186+
187+
def time_less_than_scalar1(self, dtype):
188+
(self.s < self.x)
189+
173190
def time_less_than_scalar2(self, dtype):
174-
(self.d < 1)
191+
(self.x < self.s)
175192

176193

177194
class CustomScalarFloorDivideInt(Benchmark):

‎numpy/core/code_generators/generate_umath.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -445,47 +445,47 @@ def english_upper(s):
445445
Ufunc(2, 1, None,
446446
docstrings.get('numpy.core.umath.greater'),
447447
'PyUFunc_SimpleBinaryComparisonTypeResolver',
448-
TD(all, out='?', simd=[('avx2', ints)]),
448+
TD(all, out='?', dispatch=[('loops_comparison', bints+'fd')]),
449449
[TypeDescription('O', FullTypeDescr, 'OO', 'O')],
450450
TD('O', out='?'),
451451
),
452452
'greater_equal':
453453
Ufunc(2, 1, None,
454454
docstrings.get('numpy.core.umath.greater_equal'),
455455
'PyUFunc_SimpleBinaryComparisonTypeResolver',
456-
TD(all, out='?', simd=[('avx2', ints)]),
456+
TD(all, out='?', dispatch=[('loops_comparison', bints+'fd')]),
457457
[TypeDescription('O', FullTypeDescr, 'OO', 'O')],
458458
TD('O', out='?'),
459459
),
460460
'less':
461461
Ufunc(2, 1, None,
462462
docstrings.get('numpy.core.umath.less'),
463463
'PyUFunc_SimpleBinaryComparisonTypeResolver',
464-
TD(all, out='?', simd=[('avx2', ints)]),
464+
TD(all, out='?', dispatch=[('loops_comparison', bints+'fd')]),
465465
[TypeDescription('O', FullTypeDescr, 'OO', 'O')],
466466
TD('O', out='?'),
467467
),
468468
'less_equal':
469469
Ufunc(2, 1, None,
470470
docstrings.get('numpy.core.umath.less_equal'),
471471
'PyUFunc_SimpleBinaryComparisonTypeResolver',
472-
TD(all, out='?', simd=[('avx2', ints)]),
472+
TD(all, out='?', dispatch=[('loops_comparison', bints+'fd')]),
473473
[TypeDescription('O', FullTypeDescr, 'OO', 'O')],
474474
TD('O', out='?'),
475475
),
476476
'equal':
477477
Ufunc(2, 1, None,
478478
docstrings.get('numpy.core.umath.equal'),
479479
'PyUFunc_SimpleBinaryComparisonTypeResolver',
480-
TD(all, out='?', simd=[('avx2', ints)]),
480+
TD(all, out='?', dispatch=[('loops_comparison', bints+'fd')]),
481481
[TypeDescription('O', FullTypeDescr, 'OO', 'O')],
482482
TD('O', out='?'),
483483
),
484484
'not_equal':
485485
Ufunc(2, 1, None,
486486
docstrings.get('numpy.core.umath.not_equal'),
487487
'PyUFunc_SimpleBinaryComparisonTypeResolver',
488-
TD(all, out='?', simd=[('avx2', ints)]),
488+
TD(all, out='?', dispatch=[('loops_comparison', bints+'fd')]),
489489
[TypeDescription('O', FullTypeDescr, 'OO', 'O')],
490490
TD('O', out='?'),
491491
),

‎numpy/core/setup.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1070,6 +1070,7 @@ def generate_umath_doc_header(ext, build_dir):
10701070
join('src', 'umath', 'loops_exponent_log.dispatch.c.src'),
10711071
join('src', 'umath', 'loops_hyperbolic.dispatch.c.src'),
10721072
join('src', 'umath', 'loops_modulo.dispatch.c.src'),
1073+
join('src', 'umath', 'loops_comparison.dispatch.c.src'),
10731074
join('src', 'umath', 'matmul.h.src'),
10741075
join('src', 'umath', 'matmul.c.src'),
10751076
join('src', 'umath', 'clip.h'),

‎numpy/core/src/_simd/_simd.dispatch.c.src‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,6 +472,10 @@ SIMD_IMPL_INTRIN_1(not_@bsfx@, v@bsfx@, v@bsfx@)
472472
SIMD_IMPL_INTRIN_1(tobits_@bsfx@, u64, v@bsfx@)
473473
/**end repeat**/
474474

475+
SIMD_IMPL_INTRIN_2(pack_b8_b16, vb8, vb16, vb16)
476+
SIMD_IMPL_INTRIN_4(pack_b8_b32, vb8, vb32, vb32, vb32, vb32)
477+
SIMD_IMPL_INTRIN_8(pack_b8_b64, vb8, vb64, vb64, vb64, vb64,
478+
vb64, vb64, vb64, vb64)
475479

476480
//#########################################################################
477481
//## Attach module functions
@@ -716,6 +720,11 @@ SIMD_INTRIN_DEF(not_@bsfx@)
716720
SIMD_INTRIN_DEF(tobits_@bsfx@)
717721
/**end repeat**/
718722

723+
// Pack multiple vectors into one
724+
SIMD_INTRIN_DEF(pack_b8_b16)
725+
SIMD_INTRIN_DEF(pack_b8_b32)
726+
SIMD_INTRIN_DEF(pack_b8_b64)
727+
719728
/************************************************************************/
720729
{NULL, NULL, 0, NULL}
721730
}; // PyMethodDef

‎numpy/core/src/_simd/_simd_easyintrin.inc‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,50 @@
153153
return simd_arg_to_obj(&ret); \
154154
}
155155

156+
#define SIMD_IMPL_INTRIN_8(NAME, RET, IN0, IN1, IN2, IN3, \
157+
IN4, IN5, IN6, IN7) \
158+
static PyObject *simd__intrin_##NAME \
159+
(PyObject* NPY_UNUSED(self), PyObject *args) \
160+
{ \
161+
simd_arg arg1 = {.dtype = simd_data_##IN0}; \
162+
simd_arg arg2 = {.dtype = simd_data_##IN1}; \
163+
simd_arg arg3 = {.dtype = simd_data_##IN2}; \
164+
simd_arg arg4 = {.dtype = simd_data_##IN3}; \
165+
simd_arg arg5 = {.dtype = simd_data_##IN4}; \
166+
simd_arg arg6 = {.dtype = simd_data_##IN5}; \
167+
simd_arg arg7 = {.dtype = simd_data_##IN6}; \
168+
simd_arg arg8 = {.dtype = simd_data_##IN7}; \
169+
if (!PyArg_ParseTuple( \
170+
args, "O&O&O&O&O&O&O&O&:"NPY_TOSTRING(NAME), \
171+
simd_arg_converter, &arg1, \
172+
simd_arg_converter, &arg2, \
173+
simd_arg_converter, &arg3, \
174+
simd_arg_converter, &arg4, \
175+
simd_arg_converter, &arg5, \
176+
simd_arg_converter, &arg6, \
177+
simd_arg_converter, &arg7, \
178+
simd_arg_converter, &arg8 \
179+
)) return NULL; \
180+
simd_data data = {.RET = npyv_##NAME( \
181+
arg1.data.IN0, arg2.data.IN1, \
182+
arg3.data.IN2, arg4.data.IN3, \
183+
arg5.data.IN4, arg6.data.IN5, \
184+
arg7.data.IN6, arg8.data.IN7 \
185+
)}; \
186+
simd_arg_free(&arg1); \
187+
simd_arg_free(&arg2); \
188+
simd_arg_free(&arg3); \
189+
simd_arg_free(&arg4); \
190+
simd_arg_free(&arg5); \
191+
simd_arg_free(&arg6); \
192+
simd_arg_free(&arg7); \
193+
simd_arg_free(&arg8); \
194+
simd_arg ret = { \
195+
.data = data, .dtype = simd_data_##RET \
196+
}; \
197+
return simd_arg_to_obj(&ret); \
198+
}
199+
156200
/**
157201
* Helper macros for repeating and expand a certain macro.
158202
* Mainly used for converting a scalar to an immediate constant.

‎numpy/core/src/common/simd/avx2/conversion.h‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,36 @@ NPY_FINLINE npyv_u32x2 npyv_expand_u32_u16(npyv_u16 data) {
5858
return r;
5959
}
6060

61+
// pack two 16-bit boolean into one 8-bit boolean vector
62+
NPY_FINLINE npyv_b8 npyv_pack_b8_b16(npyv_b16 a, npyv_b16 b) {
63+
__m256i ab = _mm256_packs_epi16(a, b);
64+
return npyv256_shuffle_odd(ab);
65+
}
66+
67+
// pack four 32-bit boolean vectors into one 8-bit boolean vector
68+
NPY_FINLINE npyv_b8
69+
npyv_pack_b8_b32(npyv_b32 a, npyv_b32 b, npyv_b32 c, npyv_b32 d) {
70+
__m256i ab = _mm256_packs_epi32(a, b);
71+
__m256i cd = _mm256_packs_epi32(c, d);
72+
__m256i abcd = npyv_pack_b8_b16(ab, cd);
73+
return _mm256_shuffle_epi32(abcd, _MM_SHUFFLE(3, 1, 2, 0));
74+
}
75+
76+
// pack eight 64-bit boolean vectors into one 8-bit boolean vector
77+
NPY_FINLINE npyv_b8
78+
npyv_pack_b8_b64(npyv_b64 a, npyv_b64 b, npyv_b64 c, npyv_b64 d,
79+
npyv_b64 e, npyv_b64 f, npyv_b64 g, npyv_b64 h) {
80+
__m256i ab = _mm256_packs_epi32(a, b);
81+
__m256i cd = _mm256_packs_epi32(c, d);
82+
__m256i ef = _mm256_packs_epi32(e, f);
83+
__m256i gh = _mm256_packs_epi32(g, h);
84+
__m256i abcd = _mm256_packs_epi32(ab, cd);
85+
__m256i efgh = _mm256_packs_epi32(ef, gh);
86+
__m256i all = npyv256_shuffle_odd(_mm256_packs_epi16(abcd, efgh));
87+
__m256i rev128 = _mm256_alignr_epi8(all, all, 8);
88+
return _mm256_unpacklo_epi16(all, rev128);
89+
}
90+
6191
// round to nearest integer (assuming even)
6292
#define npyv_round_s32_f32 _mm256_cvtps_epi32
6393
NPY_FINLINE npyv_s32 npyv_round_s32_f64(npyv_f64 a, npyv_f64 b)

‎numpy/core/src/common/simd/avx512/conversion.h‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,48 @@ NPY_FINLINE npyv_u32x2 npyv_expand_u32_u16(npyv_u16 data)
9090
return r;
9191
}
9292

93+
// pack two 16-bit boolean into one 8-bit boolean vector
94+
NPY_FINLINE npyv_b8 npyv_pack_b8_b16(npyv_b16 a, npyv_b16 b) {
95+
#ifdef NPY_HAVE_AVX512BW
96+
return _mm512_kunpackd((__mmask64)b, (__mmask64)a);
97+
#else
98+
const __m512i idx = _mm512_setr_epi64(0, 2, 4, 6, 1, 3, 5, 7);
99+
return _mm512_permutexvar_epi64(idx, npyv512_packs_epi16(a, b));
100+
#endif
101+
}
102+
103+
// pack four 32-bit boolean vectors into one 8-bit boolean vector
104+
NPY_FINLINE npyv_b8
105+
npyv_pack_b8_b32(npyv_b32 a, npyv_b32 b, npyv_b32 c, npyv_b32 d) {
106+
#ifdef NPY_HAVE_AVX512BW
107+
__mmask32 ab = (__mmask64)_mm512_kunpackw((__mmask32)b, (__mmask32)a);
108+
__mmask32 cd = (__mmask64)_mm512_kunpackw((__mmask32)d, (__mmask32)c);
109+
return npyv_pack_b8_b16(ab, cd);
110+
#else
111+
const __m512i idx = _mm512_setr_epi32(
112+
0, 4, 1, 5, 2, 6, 3, 7, 8, 12, 9, 13, 10, 14, 11, 15);
113+
__m256i ta = npyv512_pack_lo_hi(npyv_cvt_u32_b32(a));
114+
__m256i tb = npyv512_pack_lo_hi(npyv_cvt_u32_b32(b));
115+
__m256i tc = npyv512_pack_lo_hi(npyv_cvt_u32_b32(c));
116+
__m256i td = npyv512_pack_lo_hi(npyv_cvt_u32_b32(d));
117+
__m256i ab = _mm256_packs_epi16(ta, tb);
118+
__m256i cd = _mm256_packs_epi16(tc, td);
119+
__m512i abcd = npyv512_combine_si256(ab, cd);
120+
return _mm512_permutexvar_epi32(idx, abcd);
121+
#endif
122+
}
123+
124+
// pack eight 64-bit boolean vectors into one 8-bit boolean vector
125+
NPY_FINLINE npyv_b8
126+
npyv_pack_b8_b64(npyv_b64 a, npyv_b64 b, npyv_b64 c, npyv_b64 d,
127+
npyv_b64 e, npyv_b64 f, npyv_b64 g, npyv_b64 h) {
128+
__mmask16 ab = _mm512_kunpackb((__mmask16)b, (__mmask16)a);
129+
__mmask16 cd = _mm512_kunpackb((__mmask16)d, (__mmask16)c);
130+
__mmask16 ef = _mm512_kunpackb((__mmask16)f, (__mmask16)e);
131+
__mmask16 gh = _mm512_kunpackb((__mmask16)h, (__mmask16)g);
132+
return npyv_pack_b8_b32(ab, cd, ef, gh);
133+
}
134+
93135
// convert boolean vectors to integer bitfield
94136
NPY_FINLINE npy_uint64 npyv_tobits_b8(npyv_b8 a)
95137
{

‎numpy/core/src/common/simd/avx512/utils.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,4 +87,16 @@
8787
)); \
8888
}
8989

90+
#ifndef NPY_HAVE_AVX512BW
91+
NPYV_IMPL_AVX512_FROM_AVX2_2ARG(npyv512_packs_epi16, _mm256_packs_epi16)
92+
#else
93+
#define npyv512_packs_epi16 _mm512_packs_epi16
94+
#endif
95+
96+
NPY_FINLINE __m256i npyv512_pack_lo_hi(__m512i a) {
97+
__m256i lo = npyv512_lower_si256(a);
98+
__m256i hi = npyv512_higher_si256(a);
99+
return _mm256_packs_epi32(lo, hi);
100+
}
101+
90102
#endif // _NPY_SIMD_AVX512_UTILS_H

‎numpy/core/src/common/simd/neon/conversion.h‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,30 @@ NPY_FINLINE npyv_u32x2 npyv_expand_u32_u16(npyv_u16 data) {
8686
return r;
8787
}
8888

89+
// pack two 16-bit boolean into one 8-bit boolean vector
90+
NPY_FINLINE npyv_b8 npyv_pack_b8_b16(npyv_b16 a, npyv_b16 b) {
91+
return vcombine_u8(vmovn_u16(a), vmovn_u16(b));
92+
}
93+
94+
// pack four 32-bit boolean vectors into one 8-bit boolean vector
95+
NPY_FINLINE npyv_b8
96+
npyv_pack_b8_b32(npyv_b32 a, npyv_b32 b, npyv_b32 c, npyv_b32 d) {
97+
npyv_b16 ab = vcombine_u16(vmovn_u32(a), vmovn_u32(b));
98+
npyv_b16 cd = vcombine_u16(vmovn_u32(c), vmovn_u32(d));
99+
return npyv_pack_b8_b16(ab, cd);
100+
}
101+
102+
// pack eight 64-bit boolean vectors into one 8-bit boolean vector
103+
NPY_FINLINE npyv_b8
104+
npyv_pack_b8_b64(npyv_b64 a, npyv_b64 b, npyv_b64 c, npyv_b64 d,
105+
npyv_b64 e, npyv_b64 f, npyv_b64 g, npyv_b64 h) {
106+
npyv_b32 ab = vcombine_u32(vmovn_u64(a), vmovn_u64(b));
107+
npyv_b32 cd = vcombine_u32(vmovn_u64(c), vmovn_u64(d));
108+
npyv_b32 ef = vcombine_u32(vmovn_u64(e), vmovn_u64(f));
109+
npyv_b32 gh = vcombine_u32(vmovn_u64(g), vmovn_u64(h));
110+
return npyv_pack_b8_b32(ab, cd, ef, gh);
111+
}
112+
89113
// round to nearest integer
90114
#if NPY_SIMD_F64
91115
#define npyv_round_s32_f32 vcvtnq_s32_f32

0 commit comments

Comments
 (0)