usearch 2.15.3 → 2.16.1
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +1 -1
- package/cpp/test.cpp +16 -0
- package/include/usearch/index.hpp +21 -5
- package/include/usearch/index_dense.hpp +6 -1
- package/include/usearch/index_plugins.hpp +345 -8
- package/javascript/lib.cpp +17 -5
- package/javascript/usearch.test.js +190 -21
- package/package.json +1 -1
- package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
- package/prebuilds/linux-arm64/usearch.node +0 -0
- package/prebuilds/linux-x64/usearch.node +0 -0
- package/simsimd/.git-blame-ignore-revs +2 -0
- package/simsimd/CONTRIBUTING.md +71 -23
- package/simsimd/Cargo.lock +1 -1
- package/simsimd/Cargo.toml +4 -4
- package/simsimd/MANIFEST.in +2 -1
- package/simsimd/Package.swift +15 -3
- package/simsimd/README.md +104 -35
- package/simsimd/binding.gyp +6 -2
- package/simsimd/build.rs +2 -0
- package/simsimd/c/lib.c +180 -99
- package/simsimd/include/simsimd/binary.h +311 -102
- package/simsimd/include/simsimd/curved.h +134 -135
- package/simsimd/include/simsimd/dot.h +448 -274
- package/simsimd/include/simsimd/elementwise.h +1437 -0
- package/simsimd/include/simsimd/geospatial.h +2 -2
- package/simsimd/include/simsimd/mesh.h +2 -2
- package/simsimd/include/simsimd/probability.h +88 -88
- package/simsimd/include/simsimd/simsimd.h +1157 -772
- package/simsimd/include/simsimd/sparse.h +768 -147
- package/simsimd/include/simsimd/spatial.h +955 -376
- package/simsimd/include/simsimd/types.h +145 -71
- package/simsimd/javascript/lib.c +1 -2
- package/simsimd/package.json +4 -4
- package/simsimd/rust/lib.rs +14 -1
- package/simsimd/scripts/bench.cxx +1070 -0
- package/simsimd/scripts/bench.ipynb +187 -0
- package/simsimd/{javascript → scripts}/bench.js +1 -1
- package/simsimd/{rust/benches/sqeuclidean.rs → scripts/bench_sqeuclidean.rs} +6 -6
- package/simsimd/scripts/bench_vectors.py +700 -0
- package/simsimd/scripts/bench_vectors_live.py +213 -0
- package/simsimd/scripts/test.c +155 -0
- package/simsimd/{javascript → scripts}/test.mjs +2 -2
- package/simsimd/scripts/test.py +1410 -0
- package/simsimd/swift/{source/SimSIMD.swift → SimSIMD.swift} +6 -2
- package/simsimd/swift/{tests/SimSIMDTests.swift → Test.swift} +1 -1
- /package/simsimd/{rust/benches/cosine.rs → scripts/bench_cosine.rs} +0 -0
- /package/simsimd/{rust/benches/native.rs → scripts/bench_native.rs} +0 -0
|
@@ -6,31 +6,41 @@
|
|
|
6
6
|
*
|
|
7
7
|
* Contains:
|
|
8
8
|
* - Set Intersection ~ Jaccard Distance
|
|
9
|
+
* - Sparse Dot Products, outputting the count and weighted product
|
|
9
10
|
*
|
|
10
11
|
* For datatypes:
|
|
11
12
|
* - u16: for vocabularies under 64 thousand tokens
|
|
12
13
|
* - u32: for vocabularies under 4 billion tokens
|
|
14
|
+
* - u16 indicies + i16 weights: for weighted word counts
|
|
15
|
+
* - u16 indicies + bf16 weights: for sparse matrices
|
|
13
16
|
*
|
|
14
17
|
* For hardware architectures:
|
|
15
|
-
* - x86
|
|
16
|
-
* - Arm
|
|
18
|
+
* - x86: Ice Lake, Turin
|
|
19
|
+
* - Arm: SVE2
|
|
17
20
|
*
|
|
18
21
|
* Interestingly, to implement sparse distances and products, the most important function
|
|
19
22
|
* is analogous to `std::set_intersection`, that outputs the intersection of two sorted
|
|
20
23
|
* sequences. The naive implementation of that function would look like:
|
|
21
24
|
*
|
|
22
|
-
* std::size_t
|
|
25
|
+
* std::size_t intersection_size = 0;
|
|
23
26
|
* while (i != a_length && j != b_length) {
|
|
24
|
-
* scalar_t ai = a[i];
|
|
25
|
-
*
|
|
26
|
-
* intersection += ai == bj;
|
|
27
|
+
* scalar_t ai = a[i], bj = b[j];
|
|
28
|
+
* intersection_size += ai == bj;
|
|
27
29
|
* i += ai < bj;
|
|
28
30
|
* j += ai >= bj;
|
|
29
31
|
* }
|
|
30
32
|
*
|
|
31
33
|
* Assuming we are dealing with sparse arrays, most of the time we are just evaluating
|
|
32
34
|
* branches and skipping entries. So what if we could skip multiple entries at a time
|
|
33
|
-
* searching for the next chunk, where an intersection is possible.
|
|
35
|
+
* searching for the next chunk, where an intersection is possible. For weighted arrays:
|
|
36
|
+
*
|
|
37
|
+
* double product = 0;
|
|
38
|
+
* while (i != a_length && j != b_length) {
|
|
39
|
+
* scalar_t ai = a[i], bj = b[j];
|
|
40
|
+
* product += ai == bj ? a_weights[i] * b_weights[j] : 0;
|
|
41
|
+
* i += ai < bj;
|
|
42
|
+
* j += ai >= bj;
|
|
43
|
+
* }
|
|
34
44
|
*
|
|
35
45
|
* x86 intrinsics: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/
|
|
36
46
|
* Arm intrinsics: https://developer.arm.com/architectures/instruction-sets/intrinsics/
|
|
@@ -44,116 +54,209 @@
|
|
|
44
54
|
extern "C" {
|
|
45
55
|
#endif
|
|
46
56
|
|
|
47
|
-
// clang-format off
|
|
48
|
-
|
|
49
57
|
/* Implements the serial set intersection algorithm, similar to `std::set_intersection in C++ STL`,
|
|
50
58
|
* but uses clever galloping logic, if the arrays significantly differ in size.
|
|
51
59
|
*/
|
|
52
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u16_serial(
|
|
53
|
-
|
|
60
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_serial( //
|
|
61
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
62
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
63
|
+
simsimd_distance_t *results);
|
|
64
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_serial( //
|
|
65
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
66
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
67
|
+
simsimd_distance_t *results);
|
|
68
|
+
SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_serial( //
|
|
69
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
70
|
+
simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
|
|
71
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
72
|
+
simsimd_distance_t *results);
|
|
73
|
+
SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_serial( //
|
|
74
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
75
|
+
simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
|
|
76
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
77
|
+
simsimd_distance_t *results);
|
|
54
78
|
|
|
55
79
|
/* Implements the most naive set intersection algorithm, similar to `std::set_intersection in C++ STL`,
|
|
56
80
|
* naively enumerating the elements of two arrays.
|
|
57
81
|
*/
|
|
58
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u16_accurate(
|
|
59
|
-
|
|
82
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_accurate( //
|
|
83
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
84
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
85
|
+
simsimd_distance_t *results);
|
|
86
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_accurate( //
|
|
87
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
88
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
89
|
+
simsimd_distance_t *results);
|
|
90
|
+
SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_accurate( //
|
|
91
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
92
|
+
simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
|
|
93
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
94
|
+
simsimd_distance_t *results);
|
|
95
|
+
SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_accurate( //
|
|
96
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
97
|
+
simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
|
|
98
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
99
|
+
simsimd_distance_t *results);
|
|
60
100
|
|
|
61
101
|
/* SIMD-powered backends for Arm SVE, mostly using 32-bit arithmetic over variable-length platform-defined word sizes.
|
|
62
102
|
* Designed for Arm Graviton 3, Microsoft Cobalt, as well as Nvidia Grace and newer Ampere Altra CPUs.
|
|
63
103
|
*/
|
|
64
|
-
SIMSIMD_PUBLIC void
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
104
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_sve2( //
|
|
105
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
106
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
107
|
+
simsimd_distance_t *results);
|
|
108
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_sve2( //
|
|
109
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
110
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
111
|
+
simsimd_distance_t *results);
|
|
112
|
+
SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2( //
|
|
113
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
114
|
+
simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
|
|
115
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
116
|
+
simsimd_distance_t *results);
|
|
117
|
+
SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2( //
|
|
118
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
119
|
+
simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
|
|
120
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
121
|
+
simsimd_distance_t *results);
|
|
68
122
|
|
|
69
123
|
/* SIMD-powered backends for various generations of AVX512 CPUs.
|
|
70
124
|
* Skylake is handy, as it supports masked loads and other operations, avoiding the need for the tail loop.
|
|
125
|
+
* Ice Lake, however, is needed even for the most basic kernels to perform integer matching.
|
|
126
|
+
*/
|
|
127
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_ice( //
|
|
128
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
129
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
130
|
+
simsimd_distance_t *results);
|
|
131
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_ice( //
|
|
132
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
133
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
134
|
+
simsimd_distance_t *results);
|
|
135
|
+
|
|
136
|
+
/* SIMD-powered backends for AMD Turin CPUs with cheap VP2INTERSECT instructions.
|
|
137
|
+
* On the Intel side, only mobile Tiger Lake support them, but have prohibitively high latency.
|
|
71
138
|
*/
|
|
72
|
-
SIMSIMD_PUBLIC void
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
139
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_turin( //
|
|
140
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
141
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
142
|
+
simsimd_distance_t *results);
|
|
143
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_turin( //
|
|
144
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
145
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
146
|
+
simsimd_distance_t *results);
|
|
147
|
+
SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_turin( //
|
|
148
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
149
|
+
simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
|
|
150
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
151
|
+
simsimd_distance_t *results);
|
|
152
|
+
SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_turin( //
|
|
153
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
154
|
+
simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
|
|
155
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
156
|
+
simsimd_distance_t *results);
|
|
157
|
+
|
|
158
|
+
#define SIMSIMD_MAKE_INTERSECT_LINEAR(name, input_type, counter_type) \
|
|
159
|
+
SIMSIMD_PUBLIC void simsimd_intersect_##input_type##_##name( \
|
|
160
|
+
simsimd_##input_type##_t const *a, simsimd_##input_type##_t const *b, simsimd_size_t a_length, \
|
|
161
|
+
simsimd_size_t b_length, simsimd_distance_t *result) { \
|
|
162
|
+
simsimd_##counter_type##_t intersection_size = 0; \
|
|
163
|
+
simsimd_size_t i = 0, j = 0; \
|
|
164
|
+
while (i != a_length && j != b_length) { \
|
|
165
|
+
simsimd_##input_type##_t ai = a[i]; \
|
|
166
|
+
simsimd_##input_type##_t bj = b[j]; \
|
|
167
|
+
intersection_size += ai == bj; \
|
|
168
|
+
i += ai < bj; \
|
|
169
|
+
j += ai >= bj; \
|
|
170
|
+
} \
|
|
171
|
+
*result = intersection_size; \
|
|
90
172
|
}
|
|
91
173
|
|
|
92
174
|
SIMSIMD_MAKE_INTERSECT_LINEAR(accurate, u16, size) // simsimd_intersect_u16_accurate
|
|
93
175
|
SIMSIMD_MAKE_INTERSECT_LINEAR(accurate, u32, size) // simsimd_intersect_u32_accurate
|
|
94
176
|
|
|
95
|
-
#define
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
simsimd_size_t
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
}
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
177
|
+
#define SIMSIMD_MAKE_INTERSECT_WEIGHTED(name, variation, input_type, counter_type, weight_type, accumulator_type, \
|
|
178
|
+
load_and_convert) \
|
|
179
|
+
SIMSIMD_PUBLIC void simsimd_##variation##_##input_type##_##name( \
|
|
180
|
+
simsimd_##input_type##_t const *a, simsimd_##input_type##_t const *b, \
|
|
181
|
+
simsimd_##weight_type##_t const *a_weights, simsimd_##weight_type##_t const *b_weights, \
|
|
182
|
+
simsimd_size_t a_length, simsimd_size_t b_length, simsimd_distance_t *results) { \
|
|
183
|
+
simsimd_##counter_type##_t intersection_size = 0; \
|
|
184
|
+
simsimd_##accumulator_type##_t weights_product = 0; \
|
|
185
|
+
simsimd_size_t i = 0, j = 0; \
|
|
186
|
+
while (i != a_length && j != b_length) { \
|
|
187
|
+
simsimd_##input_type##_t ai = a[i]; \
|
|
188
|
+
simsimd_##input_type##_t bj = b[j]; \
|
|
189
|
+
int matches = ai == bj; \
|
|
190
|
+
simsimd_##counter_type##_t awi = load_and_convert(a_weights + i); \
|
|
191
|
+
simsimd_##counter_type##_t bwi = load_and_convert(b_weights + i); \
|
|
192
|
+
weights_product += matches * awi * bwi; \
|
|
193
|
+
intersection_size += matches; \
|
|
194
|
+
i += ai < bj; \
|
|
195
|
+
j += ai >= bj; \
|
|
196
|
+
} \
|
|
197
|
+
results[0] = intersection_size; \
|
|
198
|
+
results[1] = weights_product; \
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
SIMSIMD_MAKE_INTERSECT_WEIGHTED(accurate, spdot_counts, u16, size, i16, i64,
|
|
202
|
+
SIMSIMD_DEREFERENCE) // simsimd_spdot_counts_u16_accurate
|
|
203
|
+
SIMSIMD_MAKE_INTERSECT_WEIGHTED(accurate, spdot_weights, u16, size, bf16, f64,
|
|
204
|
+
SIMSIMD_BF16_TO_F32) // simsimd_spdot_weights_u16_accurate
|
|
205
|
+
|
|
206
|
+
#define SIMSIMD_MAKE_INTERSECT_GALLOPING(name, input_type, counter_type) \
|
|
207
|
+
SIMSIMD_PUBLIC simsimd_size_t simsimd_galloping_search_##input_type(simsimd_##input_type##_t const *array, \
|
|
208
|
+
simsimd_size_t start, simsimd_size_t length, \
|
|
209
|
+
simsimd_##input_type##_t val) { \
|
|
210
|
+
simsimd_size_t low = start; \
|
|
211
|
+
simsimd_size_t high = start + 1; \
|
|
212
|
+
while (high < length && array[high] < val) { \
|
|
213
|
+
low = high; \
|
|
214
|
+
high = (2 * high < length) ? 2 * high : length; \
|
|
215
|
+
} \
|
|
216
|
+
while (low < high) { \
|
|
217
|
+
simsimd_size_t mid = low + (high - low) / 2; \
|
|
218
|
+
if (array[mid] < val) { low = mid + 1; } \
|
|
219
|
+
else { high = mid; } \
|
|
220
|
+
} \
|
|
221
|
+
return low; \
|
|
222
|
+
} \
|
|
223
|
+
\
|
|
224
|
+
SIMSIMD_PUBLIC void simsimd_intersect_##input_type##_##name( \
|
|
225
|
+
simsimd_##input_type##_t const *shorter, simsimd_##input_type##_t const *longer, \
|
|
226
|
+
simsimd_size_t shorter_length, simsimd_size_t longer_length, simsimd_distance_t *result) { \
|
|
227
|
+
/* Swap arrays if necessary, as we want "longer" to be larger than "shorter" */ \
|
|
228
|
+
if (longer_length < shorter_length) { \
|
|
229
|
+
simsimd_##input_type##_t const *temp = shorter; \
|
|
230
|
+
shorter = longer; \
|
|
231
|
+
longer = temp; \
|
|
232
|
+
simsimd_size_t temp_length = shorter_length; \
|
|
233
|
+
shorter_length = longer_length; \
|
|
234
|
+
longer_length = temp_length; \
|
|
235
|
+
} \
|
|
236
|
+
\
|
|
237
|
+
/* Use the accurate implementation if galloping is not beneficial */ \
|
|
238
|
+
if (longer_length < 64 * shorter_length) { \
|
|
239
|
+
simsimd_intersect_##input_type##_accurate(shorter, longer, shorter_length, longer_length, result); \
|
|
240
|
+
return; \
|
|
241
|
+
} \
|
|
242
|
+
\
|
|
243
|
+
/* Perform galloping, shrinking the target range */ \
|
|
244
|
+
simsimd_##counter_type##_t intersection_size = 0; \
|
|
245
|
+
simsimd_size_t j = 0; \
|
|
246
|
+
for (simsimd_size_t i = 0; i < shorter_length; ++i) { \
|
|
247
|
+
simsimd_##input_type##_t shorter_i = shorter[i]; \
|
|
248
|
+
j = simsimd_galloping_search_##input_type(longer, j, longer_length, shorter_i); \
|
|
249
|
+
if (j < longer_length && longer[j] == shorter_i) { intersection_size++; } \
|
|
250
|
+
} \
|
|
251
|
+
*result = intersection_size; \
|
|
146
252
|
}
|
|
147
253
|
|
|
148
254
|
SIMSIMD_MAKE_INTERSECT_GALLOPING(serial, u16, size) // simsimd_intersect_u16_serial
|
|
149
255
|
SIMSIMD_MAKE_INTERSECT_GALLOPING(serial, u32, size) // simsimd_intersect_u32_serial
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
#pragma GCC target("avx2", "avx512f", "avx512vl", "bmi2", "lzcnt", "popcnt", "avx512bw", "avx512vbmi2")
|
|
155
|
-
#pragma clang attribute push(__attribute__((target("avx2,avx512f,avx512vl,bmi2,lzcnt,popcnt,avx512bw,avx512vbmi2"))), \
|
|
156
|
-
apply_to = function)
|
|
256
|
+
SIMSIMD_MAKE_INTERSECT_WEIGHTED(serial, spdot_counts, u16, size, i16, i32,
|
|
257
|
+
SIMSIMD_DEREFERENCE) // simsimd_spdot_counts_u16_serial
|
|
258
|
+
SIMSIMD_MAKE_INTERSECT_WEIGHTED(serial, spdot_weights, u16, size, bf16, f32,
|
|
259
|
+
SIMSIMD_BF16_TO_F32) // simsimd_spdot_weights_u16_serial
|
|
157
260
|
|
|
158
261
|
/* The AVX-512 implementations are inspired by the "Faster-Than-Native Alternatives
|
|
159
262
|
* for x86 VP2INTERSECT Instructions" paper by Guille Diez-Canas, 2022.
|
|
@@ -168,8 +271,18 @@ SIMSIMD_MAKE_INTERSECT_GALLOPING(serial, u32, size) // simsimd_intersect_u32_ser
|
|
|
168
271
|
* - `_mm512_permutexvar_epi16` - needs BW - 4-6 cycles latency
|
|
169
272
|
* - `_mm512_permutexvar_epi8` - needs VBMI - 3 cycles latency
|
|
170
273
|
*/
|
|
274
|
+
#if _SIMSIMD_TARGET_X86
|
|
275
|
+
#if SIMSIMD_TARGET_ICE
|
|
276
|
+
#pragma GCC push_options
|
|
277
|
+
#pragma GCC target("avx2", "avx512f", "avx512vl", "bmi2", "lzcnt", "popcnt", "avx512bw", "avx512vbmi2")
|
|
278
|
+
#pragma clang attribute push(__attribute__((target("avx2,avx512f,avx512vl,bmi2,lzcnt,popcnt,avx512bw,avx512vbmi2"))), \
|
|
279
|
+
apply_to = function)
|
|
171
280
|
|
|
172
|
-
|
|
281
|
+
/**
|
|
282
|
+
* @brief Analogous to `_mm512_2intersect_epi16_mask`, but compatible with Ice Lake CPUs,
|
|
283
|
+
* slightly faster than the native Tiger Lake implementation, but returns only one mask.
|
|
284
|
+
*/
|
|
285
|
+
SIMSIMD_INTERNAL simsimd_u32_t _simsimd_intersect_u16x32_ice(__m512i a, __m512i b) {
|
|
173
286
|
__m512i a1 = _mm512_alignr_epi32(a, a, 4);
|
|
174
287
|
__m512i a2 = _mm512_alignr_epi32(a, a, 8);
|
|
175
288
|
__m512i a3 = _mm512_alignr_epi32(a, a, 12);
|
|
@@ -226,7 +339,26 @@ SIMSIMD_INTERNAL simsimd_u32_t _mm512_2intersect_epi16_mask(__m512i a, __m512i b
|
|
|
226
339
|
return ~(simsimd_u32_t)(nm70 & simsimd_u32_rol(nm71, 8) & simsimd_u32_rol(nm72, 16) & simsimd_u32_ror(nm73, 8));
|
|
227
340
|
}
|
|
228
341
|
|
|
229
|
-
|
|
342
|
+
/**
|
|
343
|
+
* @brief Analogous to `_mm512_2intersect_epi32`, but compatible with Ice Lake CPUs,
|
|
344
|
+
* slightly faster than the native Tiger Lake implementation, but returns only one mask.
|
|
345
|
+
*
|
|
346
|
+
* Some latencies to keep in mind:
|
|
347
|
+
*
|
|
348
|
+
* - `_mm512_shuffle_epi32` - "VPSHUFD (ZMM, ZMM, I8)":
|
|
349
|
+
* - 1 cycle latency on Ice Lake: 1*p5
|
|
350
|
+
* - 1 cycle latency on Genoa: 1*FP123
|
|
351
|
+
* - `_mm512_mask_cmpneq_epi32_mask` - "VPCMPD (K, ZMM, ZMM, I8)":
|
|
352
|
+
* - 3 cycle latency on Ice Lake: 1*p5
|
|
353
|
+
* - 1 cycle latency on Genoa: 1*FP01
|
|
354
|
+
* - `_mm512_alignr_epi32` - "VPALIGNR (ZMM, ZMM, ZMM, I8)":
|
|
355
|
+
* - 1 cycle latency on Ice Lake: 1*p5
|
|
356
|
+
* - 2 cycle latency on Genoa: 1*FP12
|
|
357
|
+
* - `_mm512_conflict_epi32` - "VPCONFLICTD (ZMM, ZMM)":
|
|
358
|
+
* - up to 26 cycles latency on Ice Lake: 11*p0+9*p05+17*p5
|
|
359
|
+
* - up to 7 cycle latency on Genoa: 1*FP01+1*FP12
|
|
360
|
+
*/
|
|
361
|
+
SIMSIMD_INTERNAL simsimd_u16_t _simsimd_intersect_u32x16_ice(__m512i a, __m512i b) {
|
|
230
362
|
__m512i a1 = _mm512_alignr_epi32(a, a, 4);
|
|
231
363
|
__m512i b1 = _mm512_shuffle_epi32(b, _MM_PERM_ADCB);
|
|
232
364
|
__mmask16 nm00 = _mm512_cmpneq_epi32_mask(a, b);
|
|
@@ -258,8 +390,10 @@ SIMSIMD_INTERNAL simsimd_u16_t _mm512_2intersect_epi32_mask(__m512i a, __m512i b
|
|
|
258
390
|
return ~(simsimd_u16_t)(nm0 & simsimd_u16_rol(nm1, 4) & simsimd_u16_rol(nm2, 8) & simsimd_u16_ror(nm3, 4));
|
|
259
391
|
}
|
|
260
392
|
|
|
261
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u16_ice(
|
|
262
|
-
|
|
393
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_ice( //
|
|
394
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
395
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
396
|
+
simsimd_distance_t *results) {
|
|
263
397
|
|
|
264
398
|
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
265
399
|
if (a_length < 64 && b_length < 64) {
|
|
@@ -267,8 +401,8 @@ SIMSIMD_PUBLIC void simsimd_intersect_u16_ice(simsimd_u16_t const* a, simsimd_u1
|
|
|
267
401
|
return;
|
|
268
402
|
}
|
|
269
403
|
|
|
270
|
-
simsimd_u16_t const*
|
|
271
|
-
simsimd_u16_t const*
|
|
404
|
+
simsimd_u16_t const *const a_end = a + a_length;
|
|
405
|
+
simsimd_u16_t const *const b_end = b + b_length;
|
|
272
406
|
simsimd_size_t c = 0;
|
|
273
407
|
union vec_t {
|
|
274
408
|
__m512i zmm;
|
|
@@ -277,10 +411,10 @@ SIMSIMD_PUBLIC void simsimd_intersect_u16_ice(simsimd_u16_t const* a, simsimd_u1
|
|
|
277
411
|
} a_vec, b_vec;
|
|
278
412
|
|
|
279
413
|
while (a + 32 < a_end && b + 32 < b_end) {
|
|
280
|
-
a_vec.zmm = _mm512_loadu_si512((__m512i const*)a);
|
|
281
|
-
b_vec.zmm = _mm512_loadu_si512((__m512i const*)b);
|
|
414
|
+
a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
|
|
415
|
+
b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
|
|
282
416
|
|
|
283
|
-
// Intersecting registers with `
|
|
417
|
+
// Intersecting registers with `_simsimd_intersect_u16x32_ice` involves a lot of shuffling
|
|
284
418
|
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
285
419
|
simsimd_u16_t a_min;
|
|
286
420
|
simsimd_u16_t a_max = a_vec.u16[31];
|
|
@@ -290,27 +424,27 @@ SIMSIMD_PUBLIC void simsimd_intersect_u16_ice(simsimd_u16_t const* a, simsimd_u1
|
|
|
290
424
|
// If the slices don't overlap, advance the appropriate pointer
|
|
291
425
|
while (a_max < b_min && a + 64 < a_end) {
|
|
292
426
|
a += 32;
|
|
293
|
-
a_vec.zmm = _mm512_loadu_si512((__m512i const*)a);
|
|
427
|
+
a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
|
|
294
428
|
a_max = a_vec.u16[31];
|
|
295
429
|
}
|
|
296
430
|
a_min = a_vec.u16[0];
|
|
297
431
|
while (b_max < a_min && b + 64 < b_end) {
|
|
298
432
|
b += 32;
|
|
299
|
-
b_vec.zmm = _mm512_loadu_si512((__m512i const*)b);
|
|
433
|
+
b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
|
|
300
434
|
b_max = b_vec.u16[31];
|
|
301
435
|
}
|
|
302
436
|
b_min = b_vec.u16[0];
|
|
303
437
|
|
|
304
438
|
// Now we are likely to have some overlap, so we can intersect the registers
|
|
305
|
-
__mmask32 a_matches =
|
|
439
|
+
__mmask32 a_matches = _simsimd_intersect_u16x32_ice(a_vec.zmm, b_vec.zmm);
|
|
306
440
|
|
|
307
441
|
// The paper also contained a very nice procedure for exporting the matches,
|
|
308
442
|
// but we don't need it here:
|
|
309
443
|
// _mm512_mask_compressstoreu_epi16(c, a_matches, a_vec);
|
|
310
|
-
c += _mm_popcnt_u32(a_matches); //
|
|
444
|
+
c += _mm_popcnt_u32(a_matches); // MSVC has no `_popcnt32`
|
|
311
445
|
|
|
312
|
-
__m512i a_last_broadcasted = _mm512_set1_epi16(*(short const*)&a_max);
|
|
313
|
-
__m512i b_last_broadcasted = _mm512_set1_epi16(*(short const*)&b_max);
|
|
446
|
+
__m512i a_last_broadcasted = _mm512_set1_epi16(*(short const *)&a_max);
|
|
447
|
+
__m512i b_last_broadcasted = _mm512_set1_epi16(*(short const *)&b_max);
|
|
314
448
|
__mmask32 a_step_mask = _mm512_cmple_epu16_mask(a_vec.zmm, b_last_broadcasted);
|
|
315
449
|
__mmask32 b_step_mask = _mm512_cmple_epu16_mask(b_vec.zmm, a_last_broadcasted);
|
|
316
450
|
a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask);
|
|
@@ -321,8 +455,10 @@ SIMSIMD_PUBLIC void simsimd_intersect_u16_ice(simsimd_u16_t const* a, simsimd_u1
|
|
|
321
455
|
*results += c;
|
|
322
456
|
}
|
|
323
457
|
|
|
324
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u32_ice(
|
|
325
|
-
|
|
458
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_ice( //
|
|
459
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
460
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
461
|
+
simsimd_distance_t *results) {
|
|
326
462
|
|
|
327
463
|
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
328
464
|
if (a_length < 32 && b_length < 32) {
|
|
@@ -330,8 +466,8 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_ice(simsimd_u32_t const* a, simsimd_u3
|
|
|
330
466
|
return;
|
|
331
467
|
}
|
|
332
468
|
|
|
333
|
-
simsimd_u32_t const*
|
|
334
|
-
simsimd_u32_t const*
|
|
469
|
+
simsimd_u32_t const *const a_end = a + a_length;
|
|
470
|
+
simsimd_u32_t const *const b_end = b + b_length;
|
|
335
471
|
simsimd_size_t c = 0;
|
|
336
472
|
union vec_t {
|
|
337
473
|
__m512i zmm;
|
|
@@ -340,10 +476,10 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_ice(simsimd_u32_t const* a, simsimd_u3
|
|
|
340
476
|
} a_vec, b_vec;
|
|
341
477
|
|
|
342
478
|
while (a + 16 < a_end && b + 16 < b_end) {
|
|
343
|
-
a_vec.zmm = _mm512_loadu_si512((__m512i const*)a);
|
|
344
|
-
b_vec.zmm = _mm512_loadu_si512((__m512i const*)b);
|
|
479
|
+
a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
|
|
480
|
+
b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
|
|
345
481
|
|
|
346
|
-
// Intersecting registers with `
|
|
482
|
+
// Intersecting registers with `_simsimd_intersect_u32x16_ice` involves a lot of shuffling
|
|
347
483
|
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
348
484
|
simsimd_u32_t a_min;
|
|
349
485
|
simsimd_u32_t a_max = a_vec.u32[15];
|
|
@@ -353,27 +489,27 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_ice(simsimd_u32_t const* a, simsimd_u3
|
|
|
353
489
|
// If the slices don't overlap, advance the appropriate pointer
|
|
354
490
|
while (a_max < b_min && a + 32 < a_end) {
|
|
355
491
|
a += 16;
|
|
356
|
-
a_vec.zmm = _mm512_loadu_si512((__m512i const*)a);
|
|
492
|
+
a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
|
|
357
493
|
a_max = a_vec.u32[15];
|
|
358
494
|
}
|
|
359
495
|
a_min = a_vec.u32[0];
|
|
360
496
|
while (b_max < a_min && b + 32 < b_end) {
|
|
361
497
|
b += 16;
|
|
362
|
-
b_vec.zmm = _mm512_loadu_si512((__m512i const*)b);
|
|
498
|
+
b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
|
|
363
499
|
b_max = b_vec.u32[15];
|
|
364
500
|
}
|
|
365
501
|
b_min = b_vec.u32[0];
|
|
366
502
|
|
|
367
503
|
// Now we are likely to have some overlap, so we can intersect the registers
|
|
368
|
-
__mmask16 a_matches =
|
|
504
|
+
__mmask16 a_matches = _simsimd_intersect_u32x16_ice(a_vec.zmm, b_vec.zmm);
|
|
369
505
|
|
|
370
506
|
// The paper also contained a very nice procedure for exporting the matches,
|
|
371
507
|
// but we don't need it here:
|
|
372
508
|
// _mm512_mask_compressstoreu_epi32(c, a_matches, a_vec);
|
|
373
|
-
c += _mm_popcnt_u32(a_matches); //
|
|
509
|
+
c += _mm_popcnt_u32(a_matches); // MSVC has no `_popcnt32`
|
|
374
510
|
|
|
375
|
-
__m512i a_last_broadcasted = _mm512_set1_epi32(*(int const*)&a_max);
|
|
376
|
-
__m512i b_last_broadcasted = _mm512_set1_epi32(*(int const*)&b_max);
|
|
511
|
+
__m512i a_last_broadcasted = _mm512_set1_epi32(*(int const *)&a_max);
|
|
512
|
+
__m512i b_last_broadcasted = _mm512_set1_epi32(*(int const *)&b_max);
|
|
377
513
|
__mmask16 a_step_mask = _mm512_cmple_epu32_mask(a_vec.zmm, b_last_broadcasted);
|
|
378
514
|
__mmask16 b_step_mask = _mm512_cmple_epu32_mask(b_vec.zmm, a_last_broadcasted);
|
|
379
515
|
a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask);
|
|
@@ -387,9 +523,330 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_ice(simsimd_u32_t const* a, simsimd_u3
|
|
|
387
523
|
#pragma clang attribute pop
|
|
388
524
|
#pragma GCC pop_options
|
|
389
525
|
#endif // SIMSIMD_TARGET_ICE
|
|
390
|
-
#endif // SIMSIMD_TARGET_X86
|
|
391
526
|
|
|
392
|
-
#if
|
|
527
|
+
#if SIMSIMD_TARGET_TURIN
|
|
528
|
+
#pragma GCC push_options
|
|
529
|
+
#pragma GCC target("avx2", "avx512f", "avx512vl", "bmi2", "lzcnt", "popcnt", "avx512bw", "avx512vbmi2", "avx512bf16", \
|
|
530
|
+
"avx512vnni", "avx512vp2intersect")
|
|
531
|
+
#pragma clang attribute push( \
|
|
532
|
+
__attribute__((target( \
|
|
533
|
+
"avx2,avx512f,avx512vl,bmi2,lzcnt,popcnt,avx512bw,avx512vbmi2,avx512bf16,avx512vnni,avx512vp2intersect"))), \
|
|
534
|
+
apply_to = function)
|
|
535
|
+
|
|
536
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_turin( //
|
|
537
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
538
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
539
|
+
simsimd_distance_t *results) {
|
|
540
|
+
|
|
541
|
+
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
542
|
+
if (a_length < 64 && b_length < 64) {
|
|
543
|
+
simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
|
|
544
|
+
return;
|
|
545
|
+
}
|
|
546
|
+
|
|
547
|
+
//! There is no such thing as `_mm512_2intersect_epi16`, only the 32-bit variant!
|
|
548
|
+
//! So instead of jumping through 32 entries at a time, like on Ice Lake, we will
|
|
549
|
+
//! step through 16 entries at a time.
|
|
550
|
+
simsimd_u16_t const *const a_end = a + a_length;
|
|
551
|
+
simsimd_u16_t const *const b_end = b + b_length;
|
|
552
|
+
simsimd_size_t c = 0;
|
|
553
|
+
union vec_t {
|
|
554
|
+
__m256i ymm;
|
|
555
|
+
simsimd_u16_t u16[16];
|
|
556
|
+
simsimd_u8_t u8[32];
|
|
557
|
+
} a_vec, b_vec;
|
|
558
|
+
|
|
559
|
+
while (a + 16 < a_end && b + 16 < b_end) {
|
|
560
|
+
a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
|
|
561
|
+
b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
|
|
562
|
+
|
|
563
|
+
// Intersecting registers with `_mm512_2intersect_epi16_mask` involves a lot of shuffling
|
|
564
|
+
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
565
|
+
simsimd_u16_t a_min;
|
|
566
|
+
simsimd_u16_t a_max = a_vec.u16[15];
|
|
567
|
+
simsimd_u16_t b_min = b_vec.u16[0];
|
|
568
|
+
simsimd_u16_t b_max = b_vec.u16[15];
|
|
569
|
+
|
|
570
|
+
// If the slices don't overlap, advance the appropriate pointer
|
|
571
|
+
while (a_max < b_min && a + 32 < a_end) {
|
|
572
|
+
a += 16;
|
|
573
|
+
a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
|
|
574
|
+
a_max = a_vec.u16[15];
|
|
575
|
+
}
|
|
576
|
+
a_min = a_vec.u16[0];
|
|
577
|
+
while (b_max < a_min && b + 32 < b_end) {
|
|
578
|
+
b += 16;
|
|
579
|
+
b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
|
|
580
|
+
b_max = b_vec.u16[15];
|
|
581
|
+
}
|
|
582
|
+
b_min = b_vec.u16[0];
|
|
583
|
+
|
|
584
|
+
// Now we are likely to have some overlap, so we can intersect the registers
|
|
585
|
+
__m512i a_i32_vec = _mm512_cvtepu16_epi32(a_vec.ymm);
|
|
586
|
+
__m512i b_i32_vec = _mm512_cvtepu16_epi32(b_vec.ymm);
|
|
587
|
+
__mmask16 a_matches_any_in_b, b_matches_any_in_a;
|
|
588
|
+
_mm512_2intersect_epi32(a_i32_vec, b_i32_vec, &a_matches_any_in_b, &b_matches_any_in_a);
|
|
589
|
+
|
|
590
|
+
// The paper also contained a very nice procedure for exporting the matches,
|
|
591
|
+
// but we don't need it here:
|
|
592
|
+
// _mm512_mask_compressstoreu_epi16(c, a_matches_any_in_b, a_vec);
|
|
593
|
+
c += _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
|
|
594
|
+
|
|
595
|
+
__m256i a_last_broadcasted = _mm256_set1_epi16(*(short const *)&a_max);
|
|
596
|
+
__m256i b_last_broadcasted = _mm256_set1_epi16(*(short const *)&b_max);
|
|
597
|
+
__mmask16 a_step_mask = _mm256_cmple_epu16_mask(a_vec.ymm, b_last_broadcasted);
|
|
598
|
+
__mmask16 b_step_mask = _mm256_cmple_epu16_mask(b_vec.ymm, a_last_broadcasted);
|
|
599
|
+
a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask); //? Is this correct? Needs testing!
|
|
600
|
+
b += 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
|
|
601
|
+
}
|
|
602
|
+
|
|
603
|
+
simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
|
|
604
|
+
*results += c;
|
|
605
|
+
}
|
|
606
|
+
|
|
607
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_turin( //
|
|
608
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
609
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
610
|
+
simsimd_distance_t *results) {
|
|
611
|
+
|
|
612
|
+
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
613
|
+
if (a_length < 32 && b_length < 32) {
|
|
614
|
+
simsimd_intersect_u32_serial(a, b, a_length, b_length, results);
|
|
615
|
+
return;
|
|
616
|
+
}
|
|
617
|
+
|
|
618
|
+
simsimd_u32_t const *const a_end = a + a_length;
|
|
619
|
+
simsimd_u32_t const *const b_end = b + b_length;
|
|
620
|
+
simsimd_size_t c = 0;
|
|
621
|
+
union vec_t {
|
|
622
|
+
__m512i zmm;
|
|
623
|
+
simsimd_u32_t u32[16];
|
|
624
|
+
simsimd_u8_t u8[64];
|
|
625
|
+
} a_vec, b_vec;
|
|
626
|
+
|
|
627
|
+
while (a + 16 < a_end && b + 16 < b_end) {
|
|
628
|
+
a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
|
|
629
|
+
b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
|
|
630
|
+
|
|
631
|
+
// Intersecting registers with `_mm512_2intersect_epi32` involves a lot of shuffling
|
|
632
|
+
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
633
|
+
simsimd_u32_t a_min;
|
|
634
|
+
simsimd_u32_t a_max = a_vec.u32[15];
|
|
635
|
+
simsimd_u32_t b_min = b_vec.u32[0];
|
|
636
|
+
simsimd_u32_t b_max = b_vec.u32[15];
|
|
637
|
+
|
|
638
|
+
// If the slices don't overlap, advance the appropriate pointer
|
|
639
|
+
while (a_max < b_min && a + 32 < a_end) {
|
|
640
|
+
a += 16;
|
|
641
|
+
a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
|
|
642
|
+
a_max = a_vec.u32[15];
|
|
643
|
+
}
|
|
644
|
+
a_min = a_vec.u32[0];
|
|
645
|
+
while (b_max < a_min && b + 32 < b_end) {
|
|
646
|
+
b += 16;
|
|
647
|
+
b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
|
|
648
|
+
b_max = b_vec.u32[15];
|
|
649
|
+
}
|
|
650
|
+
b_min = b_vec.u32[0];
|
|
651
|
+
|
|
652
|
+
// Now we are likely to have some overlap, so we can intersect the registers
|
|
653
|
+
__mmask16 a_matches_any_in_b, b_matches_any_in_a;
|
|
654
|
+
_mm512_2intersect_epi32(a_vec.zmm, b_vec.zmm, &a_matches_any_in_b, &b_matches_any_in_a);
|
|
655
|
+
|
|
656
|
+
// The paper also contained a very nice procedure for exporting the matches,
|
|
657
|
+
// but we don't need it here:
|
|
658
|
+
// _mm512_mask_compressstoreu_epi32(c, a_matches_any_in_b, a_vec);
|
|
659
|
+
c += _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
|
|
660
|
+
|
|
661
|
+
__m512i a_last_broadcasted = _mm512_set1_epi32(*(int const *)&a_max);
|
|
662
|
+
__m512i b_last_broadcasted = _mm512_set1_epi32(*(int const *)&b_max);
|
|
663
|
+
__mmask16 a_step_mask = _mm512_cmple_epu32_mask(a_vec.zmm, b_last_broadcasted);
|
|
664
|
+
__mmask16 b_step_mask = _mm512_cmple_epu32_mask(b_vec.zmm, a_last_broadcasted);
|
|
665
|
+
a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask);
|
|
666
|
+
b += 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
|
|
667
|
+
}
|
|
668
|
+
|
|
669
|
+
simsimd_intersect_u32_serial(a, b, a_end - a, b_end - b, results);
|
|
670
|
+
*results += c;
|
|
671
|
+
}
|
|
672
|
+
|
|
673
|
+
SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_turin( //
|
|
674
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
675
|
+
simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
|
|
676
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
677
|
+
simsimd_distance_t *results) {
|
|
678
|
+
|
|
679
|
+
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
680
|
+
if (a_length < 64 && b_length < 64) {
|
|
681
|
+
simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
|
|
682
|
+
return;
|
|
683
|
+
}
|
|
684
|
+
|
|
685
|
+
//! There is no such thing as `_mm512_2intersect_epi16`, only the 32-bit variant!
|
|
686
|
+
//! So instead of jumping through 32 entries at a time, like on Ice Lake, we will
|
|
687
|
+
//! step through 16 entries at a time.
|
|
688
|
+
simsimd_u16_t const *const a_end = a + a_length;
|
|
689
|
+
simsimd_u16_t const *const b_end = b + b_length;
|
|
690
|
+
simsimd_size_t intersection_size = 0;
|
|
691
|
+
union vec_t {
|
|
692
|
+
__m256i ymm;
|
|
693
|
+
__m256 ymmps;
|
|
694
|
+
simsimd_u16_t u16[16];
|
|
695
|
+
simsimd_u8_t u8[32];
|
|
696
|
+
} a_vec, b_vec, product_vec;
|
|
697
|
+
product_vec.ymmps = _mm256_setzero_ps();
|
|
698
|
+
|
|
699
|
+
while (a + 16 < a_end && b + 16 < b_end) {
|
|
700
|
+
a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
|
|
701
|
+
b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
|
|
702
|
+
|
|
703
|
+
// Intersecting registers with `_mm512_2intersect_epi16_mask` involves a lot of shuffling
|
|
704
|
+
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
705
|
+
simsimd_u16_t a_min;
|
|
706
|
+
simsimd_u16_t a_max = a_vec.u16[15];
|
|
707
|
+
simsimd_u16_t b_min = b_vec.u16[0];
|
|
708
|
+
simsimd_u16_t b_max = b_vec.u16[15];
|
|
709
|
+
|
|
710
|
+
// If the slices don't overlap, advance the appropriate pointer
|
|
711
|
+
while (a_max < b_min && a + 32 < a_end) {
|
|
712
|
+
a += 16, a_weights += 16;
|
|
713
|
+
a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
|
|
714
|
+
a_max = a_vec.u16[15];
|
|
715
|
+
}
|
|
716
|
+
a_min = a_vec.u16[0];
|
|
717
|
+
while (b_max < a_min && b + 32 < b_end) {
|
|
718
|
+
b += 16, b_weights += 16;
|
|
719
|
+
b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
|
|
720
|
+
b_max = b_vec.u16[15];
|
|
721
|
+
}
|
|
722
|
+
b_min = b_vec.u16[0];
|
|
723
|
+
|
|
724
|
+
// Now we are likely to have some overlap, so we can intersect the registers
|
|
725
|
+
__m512i a_i32_vec = _mm512_cvtepu16_epi32(a_vec.ymm);
|
|
726
|
+
__m512i b_i32_vec = _mm512_cvtepu16_epi32(b_vec.ymm);
|
|
727
|
+
__mmask16 a_matches_any_in_b, b_matches_any_in_a;
|
|
728
|
+
_mm512_2intersect_epi32(a_i32_vec, b_i32_vec, &a_matches_any_in_b, &b_matches_any_in_a);
|
|
729
|
+
|
|
730
|
+
// The paper also contained a very nice procedure for exporting the matches,
|
|
731
|
+
// but we don't need it here:
|
|
732
|
+
// _mm512_mask_compressstoreu_epi16(intersection_size, a_matches_any_in_b, a_vec);
|
|
733
|
+
int a_matches_count_in_b = _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
|
|
734
|
+
intersection_size += a_matches_count_in_b;
|
|
735
|
+
|
|
736
|
+
// Load and shift all the relevant weights to the start of the vector before doing the dot product
|
|
737
|
+
if (a_matches_count_in_b) {
|
|
738
|
+
__m256i a_weights_vec = _mm256_lddqu_si256((__m256i const *)a_weights);
|
|
739
|
+
a_weights_vec = _mm256_maskz_compress_epi16(a_matches_any_in_b, a_weights_vec);
|
|
740
|
+
__m256i b_weights_vec = _mm256_lddqu_si256((__m256i const *)b_weights);
|
|
741
|
+
b_weights_vec = _mm256_maskz_compress_epi16(b_matches_any_in_a, b_weights_vec);
|
|
742
|
+
product_vec.ymmps = _mm256_dpbf16_ps(product_vec.ymmps, (__m256bh)a_weights_vec, (__m256bh)b_weights_vec);
|
|
743
|
+
}
|
|
744
|
+
|
|
745
|
+
__m256i a_last_broadcasted = _mm256_set1_epi16(*(short const *)&a_max);
|
|
746
|
+
__m256i b_last_broadcasted = _mm256_set1_epi16(*(short const *)&b_max);
|
|
747
|
+
__mmask16 a_step_mask = _mm256_cmple_epu16_mask(a_vec.ymm, b_last_broadcasted);
|
|
748
|
+
__mmask16 b_step_mask = _mm256_cmple_epu16_mask(b_vec.ymm, a_last_broadcasted);
|
|
749
|
+
int a_step = 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask); //? Is this correct? Needs testing!
|
|
750
|
+
int b_step = 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
|
|
751
|
+
a += a_step, a_weights += a_step;
|
|
752
|
+
b += b_step, b_weights += b_step;
|
|
753
|
+
}
|
|
754
|
+
|
|
755
|
+
simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
|
|
756
|
+
*results += intersection_size;
|
|
757
|
+
}
|
|
758
|
+
|
|
759
|
+
SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_turin( //
|
|
760
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
761
|
+
simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
|
|
762
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
763
|
+
simsimd_distance_t *results) {
|
|
764
|
+
|
|
765
|
+
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
766
|
+
if (a_length < 64 && b_length < 64) {
|
|
767
|
+
simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
|
|
768
|
+
return;
|
|
769
|
+
}
|
|
770
|
+
|
|
771
|
+
//! There is no such thing as `_mm512_2intersect_epi16`, only the 32-bit variant!
|
|
772
|
+
//! So instead of jumping through 32 entries at a time, like on Ice Lake, we will
|
|
773
|
+
//! step through 16 entries at a time.
|
|
774
|
+
simsimd_u16_t const *const a_end = a + a_length;
|
|
775
|
+
simsimd_u16_t const *const b_end = b + b_length;
|
|
776
|
+
simsimd_size_t intersection_size = 0;
|
|
777
|
+
union vec_t {
|
|
778
|
+
__m256i ymm;
|
|
779
|
+
simsimd_u16_t u16[16];
|
|
780
|
+
simsimd_u8_t u8[32];
|
|
781
|
+
} a_vec, b_vec, product_vec;
|
|
782
|
+
product_vec.ymm = _mm256_setzero_si256();
|
|
783
|
+
|
|
784
|
+
while (a + 16 < a_end && b + 16 < b_end) {
|
|
785
|
+
a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
|
|
786
|
+
b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
|
|
787
|
+
|
|
788
|
+
// Intersecting registers with `_mm512_2intersect_epi16_mask` involves a lot of shuffling
|
|
789
|
+
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
790
|
+
simsimd_u16_t a_min;
|
|
791
|
+
simsimd_u16_t a_max = a_vec.u16[15];
|
|
792
|
+
simsimd_u16_t b_min = b_vec.u16[0];
|
|
793
|
+
simsimd_u16_t b_max = b_vec.u16[15];
|
|
794
|
+
|
|
795
|
+
// If the slices don't overlap, advance the appropriate pointer
|
|
796
|
+
while (a_max < b_min && a + 32 < a_end) {
|
|
797
|
+
a += 16, a_weights += 16;
|
|
798
|
+
a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
|
|
799
|
+
a_max = a_vec.u16[15];
|
|
800
|
+
}
|
|
801
|
+
a_min = a_vec.u16[0];
|
|
802
|
+
while (b_max < a_min && b + 32 < b_end) {
|
|
803
|
+
b += 16, b_weights += 16;
|
|
804
|
+
b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
|
|
805
|
+
b_max = b_vec.u16[15];
|
|
806
|
+
}
|
|
807
|
+
b_min = b_vec.u16[0];
|
|
808
|
+
|
|
809
|
+
// Now we are likely to have some overlap, so we can intersect the registers
|
|
810
|
+
__m512i a_i32_vec = _mm512_cvtepu16_epi32(a_vec.ymm);
|
|
811
|
+
__m512i b_i32_vec = _mm512_cvtepu16_epi32(b_vec.ymm);
|
|
812
|
+
__mmask16 a_matches_any_in_b, b_matches_any_in_a;
|
|
813
|
+
_mm512_2intersect_epi32(a_i32_vec, b_i32_vec, &a_matches_any_in_b, &b_matches_any_in_a);
|
|
814
|
+
|
|
815
|
+
// The paper also contained a very nice procedure for exporting the matches,
|
|
816
|
+
// but we don't need it here:
|
|
817
|
+
// _mm512_mask_compressstoreu_epi16(intersection_size, a_matches_any_in_b, a_vec);
|
|
818
|
+
int a_matches_count_in_b = _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
|
|
819
|
+
intersection_size += a_matches_count_in_b;
|
|
820
|
+
|
|
821
|
+
// Load and shift all the relevant weights to the start of the vector before doing the dot product
|
|
822
|
+
if (a_matches_count_in_b) {
|
|
823
|
+
__m256i a_weights_vec = _mm256_lddqu_si256((__m256i const *)a_weights);
|
|
824
|
+
a_weights_vec = _mm256_maskz_compress_epi16(a_matches_any_in_b, a_weights_vec);
|
|
825
|
+
__m256i b_weights_vec = _mm256_lddqu_si256((__m256i const *)b_weights);
|
|
826
|
+
b_weights_vec = _mm256_maskz_compress_epi16(b_matches_any_in_a, b_weights_vec);
|
|
827
|
+
product_vec.ymm = _mm256_dpwssds_epi32(product_vec.ymm, a_weights_vec, b_weights_vec);
|
|
828
|
+
}
|
|
829
|
+
|
|
830
|
+
__m256i a_last_broadcasted = _mm256_set1_epi16(*(short const *)&a_max);
|
|
831
|
+
__m256i b_last_broadcasted = _mm256_set1_epi16(*(short const *)&b_max);
|
|
832
|
+
__mmask16 a_step_mask = _mm256_cmple_epu16_mask(a_vec.ymm, b_last_broadcasted);
|
|
833
|
+
__mmask16 b_step_mask = _mm256_cmple_epu16_mask(b_vec.ymm, a_last_broadcasted);
|
|
834
|
+
int a_step = 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask); //? Is this correct? Needs testing!
|
|
835
|
+
int b_step = 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
|
|
836
|
+
a += a_step, a_weights += a_step;
|
|
837
|
+
b += b_step, b_weights += b_step;
|
|
838
|
+
}
|
|
839
|
+
|
|
840
|
+
simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
|
|
841
|
+
*results += intersection_size;
|
|
842
|
+
}
|
|
843
|
+
|
|
844
|
+
#pragma clang attribute pop
|
|
845
|
+
#pragma GCC pop_options
|
|
846
|
+
#endif // SIMSIMD_TARGET_TURIN
|
|
847
|
+
#endif // _SIMSIMD_TARGET_X86
|
|
848
|
+
|
|
849
|
+
#if _SIMSIMD_TARGET_ARM
|
|
393
850
|
#if SIMSIMD_TARGET_NEON
|
|
394
851
|
#pragma GCC push_options
|
|
395
852
|
#pragma GCC target("arch=armv8.2-a")
|
|
@@ -409,8 +866,7 @@ SIMSIMD_INTERNAL int _simsimd_clz_u64(simsimd_u64_t x) {
|
|
|
409
866
|
return __builtin_clzll(x);
|
|
410
867
|
#else
|
|
411
868
|
int n = 0;
|
|
412
|
-
while ((x & 0x8000000000000000ull) == 0)
|
|
413
|
-
n++, x <<= 1;
|
|
869
|
+
while ((x & 0x8000000000000000ull) == 0) n++, x <<= 1;
|
|
414
870
|
return n;
|
|
415
871
|
#endif
|
|
416
872
|
}
|
|
@@ -448,8 +904,10 @@ SIMSIMD_INTERNAL uint16x8_t _simsimd_intersect_u16x8_neon(uint16x8_t a, uint16x8
|
|
|
448
904
|
return nm;
|
|
449
905
|
}
|
|
450
906
|
|
|
451
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u16_neon(
|
|
452
|
-
|
|
907
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_neon( //
|
|
908
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
909
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
910
|
+
simsimd_distance_t *results) {
|
|
453
911
|
|
|
454
912
|
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
455
913
|
if (a_length < 32 && b_length < 32) {
|
|
@@ -457,8 +915,8 @@ SIMSIMD_PUBLIC void simsimd_intersect_u16_neon(simsimd_u16_t const* a, simsimd_u
|
|
|
457
915
|
return;
|
|
458
916
|
}
|
|
459
917
|
|
|
460
|
-
simsimd_u16_t const*
|
|
461
|
-
simsimd_u16_t const*
|
|
918
|
+
simsimd_u16_t const *const a_end = a + a_length;
|
|
919
|
+
simsimd_u16_t const *const b_end = b + b_length;
|
|
462
920
|
union vec_t {
|
|
463
921
|
uint16x8_t u16x8;
|
|
464
922
|
simsimd_u16_t u16[8];
|
|
@@ -529,8 +987,10 @@ SIMSIMD_PUBLIC void simsimd_intersect_u16_neon(simsimd_u16_t const* a, simsimd_u
|
|
|
529
987
|
*results += vaddvq_u16(c_counts_vec.u16x8);
|
|
530
988
|
}
|
|
531
989
|
|
|
532
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u32_neon(
|
|
533
|
-
|
|
990
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_neon( //
|
|
991
|
+
simsimd_u32_t const *a, simsimd_u32_t const *b, //
|
|
992
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
993
|
+
simsimd_distance_t *results) {
|
|
534
994
|
|
|
535
995
|
// The baseline implementation for very small arrays (2 registers or less) can be quite simple:
|
|
536
996
|
if (a_length < 32 && b_length < 32) {
|
|
@@ -538,8 +998,8 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_neon(simsimd_u32_t const* a, simsimd_u
|
|
|
538
998
|
return;
|
|
539
999
|
}
|
|
540
1000
|
|
|
541
|
-
simsimd_u32_t const*
|
|
542
|
-
simsimd_u32_t const*
|
|
1001
|
+
simsimd_u32_t const *const a_end = a + a_length;
|
|
1002
|
+
simsimd_u32_t const *const b_end = b + b_length;
|
|
543
1003
|
union vec_t {
|
|
544
1004
|
uint32x4_t u32x4;
|
|
545
1005
|
simsimd_u32_t u32[4];
|
|
@@ -603,13 +1063,8 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_neon(simsimd_u32_t const* a, simsimd_u
|
|
|
603
1063
|
#pragma GCC pop_options
|
|
604
1064
|
#endif // SIMSIMD_TARGET_NEON
|
|
605
1065
|
|
|
606
|
-
|
|
607
|
-
|
|
608
|
-
#pragma GCC target("arch=armv8.2-a+sve+sve2")
|
|
609
|
-
#pragma clang attribute push(__attribute__((target("arch=armv8.2-a+sve+sve2"))), apply_to = function)
|
|
610
|
-
|
|
611
|
-
/* SVE2 introduces many new integer-oriented instructions, extending some of the NEON functionality to
|
|
612
|
-
* variable-length SVE registers. Those include "compare multiple" intrinsics:
|
|
1066
|
+
/* SVE2 introduces many new integer-oriented instructions, extending some of the NEON functionality
|
|
1067
|
+
* to variable-length SVE registers. Those include "compare multiple" intrinsics:
|
|
613
1068
|
*
|
|
614
1069
|
* - `svmatch[_u16]` that matches each scalar in first vector against all members of a 128-bit lane in the second.
|
|
615
1070
|
* - `svhistcnt[_s32]_z` does something similar, performing an inclusive prefix scan.
|
|
@@ -624,17 +1079,22 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_neon(simsimd_u32_t const* a, simsimd_u
|
|
|
624
1079
|
* https://developer.arm.com/documentation/ddi0602/2021-06/SVE-Instructions/UCLAMP--Unsigned-clamp-to-minimum-maximum-vector-?lang=en
|
|
625
1080
|
* - `TBLQ`: Table lookup quadword
|
|
626
1081
|
* https://developer.arm.com/documentation/ddi0602/2022-12/SVE-Instructions/TBLQ--Programmable-table-lookup-within-each-quadword-vector-segment--zeroing--?lang=en
|
|
627
|
-
* -
|
|
628
|
-
* https://developer.arm.com/documentation/ddi0602/2021-06/SVE-Instructions/INCB--INCD--INCH--INCW--scalar---Increment-scalar-by-multiple-of-predicate-constraint-element-count-
|
|
629
1082
|
*
|
|
630
1083
|
* Great resources for SVE2 intrinsics:
|
|
631
1084
|
*
|
|
632
1085
|
* > ARM’s Scalable Vector Extensions: A Critical Look at SVE2 For Integer Workloads
|
|
633
1086
|
* https://gist.github.com/zingaburga/805669eb891c820bd220418ee3f0d6bd
|
|
634
1087
|
*/
|
|
1088
|
+
#if SIMSIMD_TARGET_SVE2
|
|
1089
|
+
#pragma GCC push_options
|
|
1090
|
+
#pragma GCC target("arch=armv8.2-a+sve+sve2")
|
|
1091
|
+
#pragma clang attribute push(__attribute__((target("arch=armv8.2-a+sve+sve2"))), apply_to = function)
|
|
635
1092
|
|
|
636
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u16_sve2(
|
|
637
|
-
|
|
1093
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u16_sve2( //
|
|
1094
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
1095
|
+
simsimd_size_t a_length,
|
|
1096
|
+
simsimd_size_t b_length, //
|
|
1097
|
+
simsimd_distance_t *results) {
|
|
638
1098
|
|
|
639
1099
|
// A single SVE lane is 128 bits wide, so one lane fits 8 values.
|
|
640
1100
|
simsimd_size_t const register_size = svcnth();
|
|
@@ -700,8 +1160,8 @@ SIMSIMD_PUBLIC void simsimd_intersect_u16_sve2(simsimd_u16_t const* a, simsimd_u
|
|
|
700
1160
|
*results = c;
|
|
701
1161
|
}
|
|
702
1162
|
|
|
703
|
-
SIMSIMD_PUBLIC void simsimd_intersect_u32_sve2(simsimd_u32_t const*
|
|
704
|
-
simsimd_size_t b_length, simsimd_distance_t*
|
|
1163
|
+
SIMSIMD_PUBLIC void simsimd_intersect_u32_sve2(simsimd_u32_t const *a, simsimd_u32_t const *b, simsimd_size_t a_length,
|
|
1164
|
+
simsimd_size_t b_length, simsimd_distance_t *results) {
|
|
705
1165
|
|
|
706
1166
|
// A single SVE lane is 128 bits wide, so one lane fits 4 values.
|
|
707
1167
|
simsimd_size_t const register_size = svcntw();
|
|
@@ -795,10 +1255,171 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_sve2(simsimd_u32_t const* a, simsimd_u
|
|
|
795
1255
|
*results = c;
|
|
796
1256
|
}
|
|
797
1257
|
|
|
1258
|
+
SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2( //
|
|
1259
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
1260
|
+
simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
|
|
1261
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
1262
|
+
simsimd_distance_t *results) {
|
|
1263
|
+
|
|
1264
|
+
// A single SVE lane is 128 bits wide, so one lane fits 8 values.
|
|
1265
|
+
simsimd_size_t const register_size = svcnth();
|
|
1266
|
+
simsimd_size_t const lanes_count = register_size / 8;
|
|
1267
|
+
simsimd_size_t a_idx = 0, b_idx = 0;
|
|
1268
|
+
svint64_t product_vec = svdupq_n_s64(0, 0);
|
|
1269
|
+
simsimd_size_t intersection_size = 0;
|
|
1270
|
+
|
|
1271
|
+
while (a_idx < a_length && b_idx < b_length) {
|
|
1272
|
+
// Load `a_member` and broadcast it, load `b_members_vec` from memory
|
|
1273
|
+
svbool_t a_progress = svwhilelt_b16_u64(a_idx, a_length);
|
|
1274
|
+
svbool_t b_progress = svwhilelt_b16_u64(b_idx, b_length);
|
|
1275
|
+
svuint16_t a_vec = svld1_u16(a_progress, a + a_idx);
|
|
1276
|
+
svuint16_t b_vec = svld1_u16(b_progress, b + b_idx);
|
|
1277
|
+
|
|
1278
|
+
// Intersecting registers with `svmatch_u16` involves a lot of shuffling
|
|
1279
|
+
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
1280
|
+
simsimd_u16_t a_min;
|
|
1281
|
+
simsimd_u16_t a_max = svlastb(a_progress, a_vec);
|
|
1282
|
+
simsimd_u16_t b_min = svlasta(svpfalse_b(), b_vec);
|
|
1283
|
+
simsimd_u16_t b_max = svlastb(b_progress, b_vec);
|
|
1284
|
+
|
|
1285
|
+
// If the slices don't overlap, advance the appropriate pointer
|
|
1286
|
+
while (a_max < b_min && (a_idx + register_size) < a_length) {
|
|
1287
|
+
a_idx += register_size;
|
|
1288
|
+
a_progress = svwhilelt_b16_u64(a_idx, a_length);
|
|
1289
|
+
a_vec = svld1_u16(a_progress, a + a_idx);
|
|
1290
|
+
a_max = svlastb(a_progress, a_vec);
|
|
1291
|
+
}
|
|
1292
|
+
a_min = svlasta(svpfalse_b(), a_vec);
|
|
1293
|
+
while (b_max < a_min && (b_idx + register_size) < b_length) {
|
|
1294
|
+
b_idx += register_size;
|
|
1295
|
+
b_progress = svwhilelt_b16_u64(b_idx, b_length);
|
|
1296
|
+
b_vec = svld1_u16(b_progress, b + b_idx);
|
|
1297
|
+
b_max = svlastb(b_progress, b_vec);
|
|
1298
|
+
}
|
|
1299
|
+
b_min = svlasta(svpfalse_b(), b_vec);
|
|
1300
|
+
|
|
1301
|
+
// Before we evaluate the intersection size, obfurscating the order in `b_vec`,
|
|
1302
|
+
// let's estimate how much we will need to advance the pointers afterwards.
|
|
1303
|
+
// For that, we don't even need to broadcast the values in SVE, as the whole
|
|
1304
|
+
// register can be compared against a scalar:
|
|
1305
|
+
//
|
|
1306
|
+
// svuint16_t a_last_broadcasted = svdup_n_u16(a_max);
|
|
1307
|
+
// svuint16_t b_last_broadcasted = svdup_n_u16(b_max);
|
|
1308
|
+
svbool_t a_mask = svcmple_n_u16(a_progress, a_vec, b_max);
|
|
1309
|
+
svbool_t b_mask = svcmple_n_u16(b_progress, b_vec, a_max);
|
|
1310
|
+
simsimd_u64_t a_step = svcntp_b16(a_progress, a_mask);
|
|
1311
|
+
simsimd_u64_t b_step = svcntp_b16(b_progress, b_mask);
|
|
1312
|
+
|
|
1313
|
+
// Compare `a_vec` with each lane of `b_vec`
|
|
1314
|
+
svint16_t a_weights_vec = svld1_s16(a_progress, a_weights + a_idx);
|
|
1315
|
+
svint16_t b_weights_vec = svld1_s16(b_progress, b_weights + b_idx);
|
|
1316
|
+
for (simsimd_size_t i = 0; i < lanes_count; i++) {
|
|
1317
|
+
svbool_t equal_mask = svmatch_u16(a_progress, a_vec, b_vec);
|
|
1318
|
+
svint16_t b_equal_weights_vec = svsel_s16(equal_mask, b_weights_vec, svdup_n_s16(0.f));
|
|
1319
|
+
product_vec = svdot_s64(product_vec, a_weights_vec, b_equal_weights_vec);
|
|
1320
|
+
b_vec = svext_u16(b_vec, b_vec, 8);
|
|
1321
|
+
intersection_size += svcntp_b16(svptrue_b16(), equal_mask);
|
|
1322
|
+
}
|
|
1323
|
+
|
|
1324
|
+
// Advance
|
|
1325
|
+
a_idx += a_step;
|
|
1326
|
+
b_idx += b_step;
|
|
1327
|
+
}
|
|
1328
|
+
results[0] = (simsimd_distance_t)intersection_size;
|
|
1329
|
+
results[1] = svaddv_s64(svptrue_b64(), product_vec);
|
|
1330
|
+
}
|
|
1331
|
+
|
|
798
1332
|
#pragma clang attribute pop
|
|
799
1333
|
#pragma GCC pop_options
|
|
800
1334
|
#endif // SIMSIMD_TARGET_SVE2
|
|
801
|
-
|
|
1335
|
+
|
|
1336
|
+
#if SIMSIMD_TARGET_SVE2 && SIMSIMD_TARGET_SVE_BF16
|
|
1337
|
+
#pragma GCC push_options
|
|
1338
|
+
#pragma GCC target("arch=armv8.6-a+sve+sve2+bf16")
|
|
1339
|
+
#pragma clang attribute push(__attribute__((target("arch=armv8.6-a+sve+sve2+bf16"))), apply_to = function)
|
|
1340
|
+
|
|
1341
|
+
SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2( //
|
|
1342
|
+
simsimd_u16_t const *a, simsimd_u16_t const *b, //
|
|
1343
|
+
simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
|
|
1344
|
+
simsimd_size_t a_length, simsimd_size_t b_length, //
|
|
1345
|
+
simsimd_distance_t *results) {
|
|
1346
|
+
|
|
1347
|
+
// A single SVE lane is 128 bits wide, so one lane fits 8 values.
|
|
1348
|
+
simsimd_size_t const register_size = svcnth();
|
|
1349
|
+
simsimd_size_t const lanes_count = register_size / 8;
|
|
1350
|
+
simsimd_size_t a_idx = 0, b_idx = 0;
|
|
1351
|
+
svfloat32_t product_vec = svdupq_n_f32(0.f, 0.f, 0.f, 0.f);
|
|
1352
|
+
simsimd_size_t intersection_size = 0;
|
|
1353
|
+
|
|
1354
|
+
while (a_idx < a_length && b_idx < b_length) {
|
|
1355
|
+
// Load `a_member` and broadcast it, load `b_members_vec` from memory
|
|
1356
|
+
svbool_t a_progress = svwhilelt_b16_u64(a_idx, a_length);
|
|
1357
|
+
svbool_t b_progress = svwhilelt_b16_u64(b_idx, b_length);
|
|
1358
|
+
svuint16_t a_vec = svld1_u16(a_progress, a + a_idx);
|
|
1359
|
+
svuint16_t b_vec = svld1_u16(b_progress, b + b_idx);
|
|
1360
|
+
|
|
1361
|
+
// Intersecting registers with `svmatch_u16` involves a lot of shuffling
|
|
1362
|
+
// and comparisons, so we want to avoid it if the slices don't overlap at all..
|
|
1363
|
+
simsimd_u16_t a_min;
|
|
1364
|
+
simsimd_u16_t a_max = svlastb(a_progress, a_vec);
|
|
1365
|
+
simsimd_u16_t b_min = svlasta(svpfalse_b(), b_vec);
|
|
1366
|
+
simsimd_u16_t b_max = svlastb(b_progress, b_vec);
|
|
1367
|
+
|
|
1368
|
+
// If the slices don't overlap, advance the appropriate pointer
|
|
1369
|
+
while (a_max < b_min && (a_idx + register_size) < a_length) {
|
|
1370
|
+
a_idx += register_size;
|
|
1371
|
+
a_progress = svwhilelt_b16_u64(a_idx, a_length);
|
|
1372
|
+
a_vec = svld1_u16(a_progress, a + a_idx);
|
|
1373
|
+
a_max = svlastb(a_progress, a_vec);
|
|
1374
|
+
}
|
|
1375
|
+
a_min = svlasta(svpfalse_b(), a_vec);
|
|
1376
|
+
while (b_max < a_min && (b_idx + register_size) < b_length) {
|
|
1377
|
+
b_idx += register_size;
|
|
1378
|
+
b_progress = svwhilelt_b16_u64(b_idx, b_length);
|
|
1379
|
+
b_vec = svld1_u16(b_progress, b + b_idx);
|
|
1380
|
+
b_max = svlastb(b_progress, b_vec);
|
|
1381
|
+
}
|
|
1382
|
+
b_min = svlasta(svpfalse_b(), b_vec);
|
|
1383
|
+
|
|
1384
|
+
// Before we evaluate the intersection size, obfurscating the order in `b_vec`,
|
|
1385
|
+
// let's estimate how much we will need to advance the pointers afterwards.
|
|
1386
|
+
// For that, we don't even need to broadcast the values in SVE, as the whole
|
|
1387
|
+
// register can be compared against a scalar:
|
|
1388
|
+
//
|
|
1389
|
+
// svuint16_t a_last_broadcasted = svdup_n_u16(a_max);
|
|
1390
|
+
// svuint16_t b_last_broadcasted = svdup_n_u16(b_max);
|
|
1391
|
+
svbool_t a_mask = svcmple_n_u16(a_progress, a_vec, b_max);
|
|
1392
|
+
svbool_t b_mask = svcmple_n_u16(b_progress, b_vec, a_max);
|
|
1393
|
+
simsimd_u64_t a_step = svcntp_b16(a_progress, a_mask);
|
|
1394
|
+
simsimd_u64_t b_step = svcntp_b16(b_progress, b_mask);
|
|
1395
|
+
|
|
1396
|
+
// Compare `a_vec` with each lane of `b_vec`
|
|
1397
|
+
svbfloat16_t a_weights_vec = svld1_bf16(a_progress, (__bf16 const *)a_weights + a_idx);
|
|
1398
|
+
svbfloat16_t b_weights_vec = svld1_bf16(b_progress, (__bf16 const *)b_weights + b_idx);
|
|
1399
|
+
for (simsimd_size_t i = 0; i < lanes_count; i++) {
|
|
1400
|
+
svbool_t equal_mask = svmatch_u16(a_progress, a_vec, b_vec);
|
|
1401
|
+
//! The `svsel_bf16` intrinsic is broken in many compilers, not returning the correct type.
|
|
1402
|
+
//! So we reinterprete floats as integers and apply `svsel_s16`, but the `svreinterpret_s16_bs16`
|
|
1403
|
+
//! and `svreinterpret_bf16_s16` are not always properly defined!
|
|
1404
|
+
svint16_t b_equal_weights_vec =
|
|
1405
|
+
svsel_s16(equal_mask, svreinterpret_s16_bf16(b_weights_vec), svdup_n_s16(0));
|
|
1406
|
+
product_vec = svbfdot_f32(product_vec, a_weights_vec, svreinterpret_bf16_s16(b_equal_weights_vec));
|
|
1407
|
+
b_vec = svext_u16(b_vec, b_vec, 8);
|
|
1408
|
+
intersection_size += svcntp_b16(svptrue_b16(), equal_mask);
|
|
1409
|
+
}
|
|
1410
|
+
|
|
1411
|
+
// Advance
|
|
1412
|
+
a_idx += a_step;
|
|
1413
|
+
b_idx += b_step;
|
|
1414
|
+
}
|
|
1415
|
+
results[0] = (simsimd_distance_t)intersection_size;
|
|
1416
|
+
results[1] = svaddv_f32(svptrue_b32(), product_vec);
|
|
1417
|
+
}
|
|
1418
|
+
|
|
1419
|
+
#pragma clang attribute pop
|
|
1420
|
+
#pragma GCC pop_options
|
|
1421
|
+
#endif // SIMSIMD_TARGET_SVE2 && SIMSIMD_TARGET_SVE_BF16
|
|
1422
|
+
#endif // _SIMSIMD_TARGET_ARM
|
|
802
1423
|
|
|
803
1424
|
#ifdef __cplusplus
|
|
804
1425
|
}
|